"""Import API client for bulk/legacy data ingestion.

Provides a high-level interface for importing historical telemetry data
into fishinfo. Superadmin access required.

Example - Batch import from an old database:
    ```python
    import asyncio
    from datetime import datetime, timezone
    from uuid import UUID

    from fishinfo import FishinfoConfig
    from fishinfo.api import ImportApi

    async def migrate_legacy_data():
        config = FishinfoConfig(
            base_url="https://fishinfo.api.vikdata.local:8080",
            api_key="vk_live_your_superadmin_key",
        )

        async with FishinfoClient(config) as client:
            import_api = ImportApi(client)

            # Get the expected schema
            schema = await import_api.get_schema()
            print(f"Supported data types: {schema.data_types}")

            # Import records in batches
            records = [
                BulkTelemetryRecord(
                    time=datetime(2020, 6, 15, 14, 30, tzinfo=timezone.utc),
                    vessel_id=UUID("550e8400-e29b-41d4-a716-446655440000"),
                    data_type="position",
                    data={"latitude": 64.1466, "longitude": -21.9426, "speed": 12.5},
                ),
                # ... more records
            ]

            result = await import_api.batch_import(
                tenant_id=UUID("your-tenant-id"),
                records=records,
            )
            print(f"Imported: {result.rows_inserted} rows")
            print(f"  Hot storage: {result.rows_to_hot}")
            print(f"  Cold storage: {result.rows_to_cold}")
            print(f"  Failed: {result.rows_failed}")

    asyncio.run(migrate_legacy_data())
    ```

Example - Large S3 upload job:
    ```python
    async def migrate_large_dataset():
        async with FishinfoClient(config) as client:
            import_api = ImportApi(client)

            # Create a job with 5 parts
            response = await import_api.create_job(
                tenant_id=UUID("your-tenant-id"),
                name="Historical GPS 2018-2022",
                total_parts=5,
                format="jsonl",
            )
            job = response.job
            parts = response.parts

            # Upload JSONL files to S3 using pre-signed URLs
            for part in parts:
                if part.upload_url:
                    file_path = f"data/part_{part.part_number}.jsonl"
                    await import_api.upload_file(job.id, part, file_path)

            # Start processing
            job = await import_api.start_processing(job.id)

            # Poll for completion
            while job.status in ("processing", "uploading"):
                await asyncio.sleep(5)
                job = await import_api.get_job(job.id)

            print(f"Job {job.status}: {job.total_rows} rows imported")
    ```
"""

import hashlib
import json
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING
from uuid import UUID

from fishinfo.models.import_job import (
    BulkImportResponse,
    BulkTelemetryRecord,
    CreateImportJobResponse,
    ImportJob,
    ImportJobListResponse,
    ImportJobPart,
    ImportSchemaResponse,
)

if TYPE_CHECKING:
    from fishinfo.client import FishinfoClient


class ImportApi:
    """API client for bulk/legacy data import operations.

    All operations require superadmin access. This client wraps the
    fishinfo import REST API and provides convenience methods for
    common migration workflows.

    Attributes:
        client: FishinfoClient instance
    """

    def __init__(self, client: "FishinfoClient") -> None:
        self.client = client

    # =========================================================================
    # Job Management
    # =========================================================================

    async def create_job(
        self,
        tenant_id: UUID,
        name: str,
        total_parts: int,
        format: str = "jsonl",
        description: str | None = None,
        enable_domain_processing: bool = False,
    ) -> CreateImportJobResponse:
        """Create a new import job with pre-signed S3 upload URLs.

        Args:
            tenant_id: Target tenant for the import
            name: Human-readable name (e.g., "Historical GPS 2018-2022")
            total_parts: Number of files that will be uploaded
            format: Data format - "jsonl" or "parquet" (default: "jsonl")
            description: Optional description
            enable_domain_processing: Whether to trigger trip/zone events (default: False)

        Returns:
            CreateImportJobResponse with job details and parts with upload URLs
        """
        body = {
            "tenant_id": str(tenant_id),
            "name": name,
            "total_parts": total_parts,
            "format": format,
            "enable_domain_processing": enable_domain_processing,
        }
        if description:
            body["description"] = description

        data = await self.client._request("POST", "/import/jobs", json=body)
        return CreateImportJobResponse(**data)

    async def get_job(self, job_id: UUID) -> ImportJob:
        """Get import job details.

        Args:
            job_id: Job ID

        Returns:
            ImportJob with current status and statistics
        """
        data = await self.client._request("GET", f"/import/jobs/{job_id}")
        return ImportJob(**data)

    async def list_jobs(
        self,
        tenant_id: UUID | None = None,
        status: str | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> ImportJobListResponse:
        """List import jobs.

        Args:
            tenant_id: Filter by tenant (optional)
            status: Filter by status (optional)
            limit: Page size (default: 50)
            offset: Pagination offset (default: 0)

        Returns:
            ImportJobListResponse with paginated results
        """
        params: dict = {"limit": limit, "offset": offset}
        if tenant_id:
            params["tenant_id"] = str(tenant_id)
        if status:
            params["status"] = status

        data = await self.client._request("GET", "/import/jobs", params=params)
        return ImportJobListResponse(**data)

    async def get_parts(self, job_id: UUID) -> list[ImportJobPart]:
        """Get all parts for an import job.

        Args:
            job_id: Job ID

        Returns:
            List of ImportJobPart with upload status
        """
        data = await self.client._request("GET", f"/import/jobs/{job_id}/parts")
        return [ImportJobPart(**p) for p in data]

    async def cancel_job(self, job_id: UUID) -> ImportJob:
        """Cancel an import job.

        Args:
            job_id: Job ID

        Returns:
            Updated ImportJob
        """
        data = await self.client._request("POST", f"/import/jobs/{job_id}/cancel")
        return ImportJob(**data)

    # =========================================================================
    # Upload Management
    # =========================================================================

    async def confirm_part_upload(
        self,
        job_id: UUID,
        part_number: int,
        size_bytes: int,
        checksum: str | None = None,
    ) -> ImportJobPart:
        """Confirm that a part has been uploaded to S3.

        Call this after uploading a file via the pre-signed URL.

        Args:
            job_id: Job ID
            part_number: Part number
            size_bytes: Size of the uploaded file
            checksum: Optional SHA256 checksum

        Returns:
            Updated ImportJobPart
        """
        body: dict = {"size_bytes": size_bytes}
        if checksum:
            body["checksum"] = checksum

        data = await self.client._request(
            "POST",
            f"/import/jobs/{job_id}/parts/{part_number}/confirm",
            json=body,
        )
        return ImportJobPart(**data)

    async def refresh_urls(self, job_id: UUID) -> list[ImportJobPart]:
        """Refresh expired pre-signed upload URLs.

        Args:
            job_id: Job ID

        Returns:
            Parts with refreshed upload URLs
        """
        data = await self.client._request(
            "POST", f"/import/jobs/{job_id}/parts/refresh-urls"
        )
        return [ImportJobPart(**p) for p in data]

    async def upload_file(
        self,
        job_id: UUID,
        part: ImportJobPart,
        file_path: str | Path,
    ) -> ImportJobPart:
        """Upload a file to S3 using the part's pre-signed URL, then confirm.

        This is a convenience method that handles the full upload workflow:
        1. Reads the file
        2. Computes SHA256 checksum
        3. Uploads to S3 via the pre-signed URL
        4. Confirms the upload with the fishinfo API

        Args:
            job_id: Job ID
            part: ImportJobPart with a valid upload_url
            file_path: Local path to the file to upload

        Returns:
            Updated ImportJobPart after confirmation

        Raises:
            ValueError: If part has no upload_url
            FileNotFoundError: If file doesn't exist
        """
        import httpx

        if not part.upload_url:
            raise ValueError(
                f"Part {part.part_number} has no upload URL. "
                "Call refresh_urls() to get new URLs."
            )

        path = Path(file_path)
        if not path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        # Compute checksum
        file_hash = hashlib.sha256()
        with open(path, "rb") as f:
            for chunk in iter(lambda: f.read(8192), b""):
                file_hash.update(chunk)
        checksum = file_hash.hexdigest()
        size_bytes = path.stat().st_size

        # Upload to S3 via pre-signed URL
        content_type = (
            "application/x-ndjson"
            if path.suffix == ".jsonl"
            else "application/octet-stream"
        )

        async with httpx.AsyncClient() as http:
            with open(path, "rb") as f:
                response = await http.put(
                    part.upload_url,
                    content=f.read(),
                    headers={"Content-Type": content_type},
                    timeout=600.0,  # 10 min timeout for large files
                )
                response.raise_for_status()

        # Confirm upload
        return await self.confirm_part_upload(
            job_id=job_id,
            part_number=part.part_number,
            size_bytes=size_bytes,
            checksum=checksum,
        )

    # =========================================================================
    # Processing
    # =========================================================================

    async def start_processing(self, job_id: UUID) -> ImportJob:
        """Start processing all uploaded parts.

        Triggers background processing on the server. Returns immediately.
        Poll get_job() to check progress.

        Args:
            job_id: Job ID

        Returns:
            ImportJob with status='processing'
        """
        data = await self.client._request(
            "POST", f"/import/jobs/{job_id}/process"
        )
        return ImportJob(**data)

    async def wait_for_completion(
        self,
        job_id: UUID,
        poll_interval: float = 5.0,
        timeout: float | None = None,
    ) -> ImportJob:
        """Wait for an import job to complete.

        Polls the job status until it reaches a terminal state
        (completed, failed, or cancelled).

        Args:
            job_id: Job ID
            poll_interval: Seconds between status checks (default: 5)
            timeout: Maximum seconds to wait (default: None = no timeout)

        Returns:
            ImportJob in terminal state

        Raises:
            TimeoutError: If timeout is reached before completion
        """
        import asyncio
        import time

        start = time.monotonic()
        while True:
            job = await self.get_job(job_id)
            if job.status in ("completed", "failed", "cancelled"):
                return job

            if timeout and (time.monotonic() - start) > timeout:
                raise TimeoutError(
                    f"Import job {job_id} did not complete within {timeout}s"
                )

            await asyncio.sleep(poll_interval)

    # =========================================================================
    # Batch Import (direct API)
    # =========================================================================

    async def batch_import(
        self,
        tenant_id: UUID,
        records: list[BulkTelemetryRecord],
    ) -> BulkImportResponse:
        """Import a batch of records directly (max 10,000 per call).

        For smaller imports that don't need the full S3 upload workflow.
        Records are validated and routed to hot/cold storage immediately.

        Args:
            tenant_id: Target tenant
            records: List of telemetry records (max 10,000)

        Returns:
            BulkImportResponse with insertion statistics
        """
        body = {
            "records": [r.model_dump(mode="json") for r in records],
        }

        data = await self.client._request(
            "POST", f"/import/batch/{tenant_id}", json=body
        )
        return BulkImportResponse(**data)

    async def batch_import_iter(
        self,
        tenant_id: UUID,
        records: list[BulkTelemetryRecord],
        batch_size: int = 5000,
    ) -> BulkImportResponse:
        """Import records in automatic batches.

        Splits a large list into batches of `batch_size` and imports
        each batch sequentially. Aggregates the results.

        Args:
            tenant_id: Target tenant
            records: List of telemetry records (any size)
            batch_size: Records per batch (default: 5000, max: 10000)

        Returns:
            Aggregated BulkImportResponse across all batches
        """
        batch_size = min(batch_size, 10000)
        total_inserted = 0
        total_hot = 0
        total_cold = 0
        total_failed = 0
        all_errors: list[dict] = []

        for i in range(0, len(records), batch_size):
            batch = records[i : i + batch_size]
            result = await self.batch_import(tenant_id, batch)
            total_inserted += result.rows_inserted
            total_hot += result.rows_to_hot
            total_cold += result.rows_to_cold
            total_failed += result.rows_failed
            if result.errors:
                # Offset error indices by batch position
                for error in result.errors:
                    adjusted = {**error}
                    if "index" in adjusted:
                        adjusted["index"] = adjusted["index"] + i
                    all_errors.append(adjusted)

        return BulkImportResponse(
            rows_inserted=total_inserted,
            rows_to_hot=total_hot,
            rows_to_cold=total_cold,
            rows_failed=total_failed,
            errors=all_errors[:100] if all_errors else None,
        )

    # =========================================================================
    # Schema / Template
    # =========================================================================

    async def get_schema(self) -> ImportSchemaResponse:
        """Get the expected data format/schema for import files.

        Returns format descriptions, supported data types, and example records.

        Returns:
            ImportSchemaResponse with format info and examples
        """
        data = await self.client._request("GET", "/import/schema")
        return ImportSchemaResponse(**data)

    # =========================================================================
    # JSONL File Generation Helper
    # =========================================================================

    @staticmethod
    def write_jsonl(
        records: list[BulkTelemetryRecord],
        output_path: str | Path,
    ) -> int:
        """Write records to a JSONL file for S3 upload.

        Convenience method for generating properly formatted JSONL files
        from BulkTelemetryRecord objects.

        Args:
            records: List of telemetry records
            output_path: Path to write the JSONL file

        Returns:
            Number of records written
        """
        path = Path(output_path)
        path.parent.mkdir(parents=True, exist_ok=True)

        with open(path, "w") as f:
            for record in records:
                line = record.model_dump_json()
                f.write(line + "\n")

        return len(records)
