from __future__ import annotations

from pathlib import Path
import time
import uuid
from typing import Any, Dict, List, Optional, Union

from specific_ai.platform.base_client import BaseClient
from specific_ai.platform.exceptions import APIErrorDetails, NetworkError, NotFoundError
from specific_ai.platform.schemas import ModelMetrics, ModelResponse


class ModelManager:
    """Interact with trained model artifacts and deployments."""

    def __init__(self, client: BaseClient, *, client_id: str):
        self._client = client
        self._client_id = client_id

    def get_metrics(
        self,
        *,
        task_id: str,
        version: Optional[float] = None,
        comparison_mode: bool = False,
        compare_with_version: Optional[float] = None,
        compare_with_file: bool = False,
        sub_task_type: Optional[str] = None,
        task_type: Optional[str] = None,
    ) -> ModelMetrics:
        """
        Fetch evaluation metrics for a model.

        Args:
            task_id: Task id.
            version: Optional version to fetch; defaults to latest.
            comparison_mode: Whether to compare two student versions.
            compare_with_version: Version to compare against when `comparison_mode=True`.
            compare_with_file: Whether to compare against an uploaded file baseline.
            sub_task_type: Optional sub-task type.
            task_type: Optional task type override.

        Returns:
            ModelMetrics: Evaluation response (contains raw extra fields too).

        Raises:
            SpecificAIAPIError: If the API returns an error.
        """
        payload: Dict[str, Any] = {
            "client_id": self._client_id,
            "llm_usecase_id": task_id,
            "version": version,
            "comparison_mode": comparison_mode,
            "compare_with_version": compare_with_version,
            "compare_with_file": compare_with_file,
            "sub_task_type": sub_task_type,
            "task_type": task_type,
        }
        data = self._client.request_json("POST", "/evaluation", json_body=payload)
        return ModelMetrics.parse_obj(data)

    def list_versions(self, *, task_id: str) -> List[Dict[str, Any]]:
        """
        List available model versions for a task.

        Under the hood, this calls `POST /trainings` and filters by `llm_usecase_id == task_id`.
        It also tries to attach an `accuracy` field per version (best-effort) by calling
        `get_metrics(task_id=..., version=...)`. Accuracy may be `None` if the task type does not
        report it, or if metrics are not available for that version.

        Args:
            task_id: Task id (llm_usecase_id).

        Returns:
            list[dict]: One item per training/version, sorted by `version` descending. Each item
            contains:
                - version: float | None
                - created_at: str | None
                - finished_at: str | None
                - status: str | None
                - progress: int | None
                - distillation_event_id: str | None
                - training_id: str | None
                - created_by_user: dict | None
                - accuracy: float | None

        Raises:
            SpecificAIAPIError: If the API returns an error.
        """
        data = self._client.request_json(
            "POST", "/trainings", json_body={"value": self._client_id}
        )
        trainings = []
        if isinstance(data, dict):
            trainings = data.get("trainings") or []

        out: List[Dict[str, Any]] = []
        for t in trainings:
            if not isinstance(t, dict):
                continue
            if t.get("llm_usecase_id") != task_id:
                continue
            version = t.get("version")
            metrics: Optional[ModelMetrics] = None
            if isinstance(version, (int, float)):
                try:
                    metrics = self.get_metrics(task_id=task_id, version=float(version))
                except Exception:
                    metrics = None

            out.append(
                {
                    "version": version,
                    "created_at": t.get("created_datetime"),
                    "finished_at": t.get("finished_datetime"),
                    "status": t.get("status"),
                    "progress": t.get("progress"),
                    "distillation_event_id": t.get("distillation_event_id"),
                    "training_id": t.get("id"),
                    "created_by_user": t.get("created_by_user"),
                    "metrics": metrics,
                }
            )

        out.sort(
            key=lambda x: (
                float(x["version"])
                if isinstance(x.get("version"), (int, float))
                else -1
            ),
            reverse=True,
        )
        return out

    def delete_version(self, *, task_id: str, version: float) -> Dict[str, Any]:
        """
        Delete a specific model version (training run).

        This is intended for cleanup of old/bad experiments. The SDK will look up the training's
        `distillation_event_id` for the requested `task_id` + `version`, then call
        `POST /delete_training`.

        Args:
            task_id: Task id (llm_usecase_id).
            version: Version number to delete.

        Returns:
            dict: API response (typically {"success": bool, "message": str}).

        Raises:
            NotFoundError: If the requested version does not exist for this task.
            SpecificAIAPIError: If the API returns an error.
        """
        versions = self.list_versions(task_id=task_id)
        match = next(
            (
                v
                for v in versions
                if isinstance(v.get("version"), (int, float))
                and float(v["version"]) == float(version)
            ),
            None,
        )
        if not match or not match.get("distillation_event_id"):
            raise NotFoundError(
                APIErrorDetails(
                    status_code=404,
                    message=f"Training version not found (task_id={task_id}, version={version})",
                    response_json={"task_id": task_id, "version": version},
                )
            )

        return self._client.request_json(
            "POST",
            "/delete_training",
            json_body={
                "client_id": self._client_id,
                "distillation_event_id": match["distillation_event_id"],
            },
        )  # type: ignore[return-value]

    def deploy(
        self,
        *,
        task_id: str,
        task_type: str,
        deployed_version: float,
        request_type: str = "deploy",
    ) -> Dict[str, Any]:
        """
        Deploy or undeploy a model version on the inference infrastructure.

        Args:
            task_id: Task id.
            task_type: Task type string (backend enum value).
            deployed_version: Version to deploy.
            request_type: "deploy" or "undeploy".

        Returns:
            dict: API response (typically {"success": bool}).

        Raises:
            SpecificAIAPIError: If the API returns an error.
        """
        payload = {
            "client_id": self._client_id,
            "llm_usecase_id": task_id,
            "task_type": task_type,
            "request_type": request_type,
            "deployed_version": deployed_version,
        }
        data = self._client.request_json("POST", "/model-deployment", json_body=payload)
        return data  # type: ignore[return-value]

    def evaluate(
        self,
        *,
        task_id: str,
        version: float | None = None,
    ) -> Dict[str, Any]:
        """
        Trigger re-evaluation for an existing model version.

        Under the hood this calls `POST /recalculate-evaluation`.
        If `version` is omitted, the SDK defaults to the latest available version from metrics.
        """
        if version is None:
            metrics = self.get_metrics(task_id=task_id)
            if not metrics.versions:
                raise NotFoundError(
                    APIErrorDetails(
                        status_code=404,
                        message=f"No versions found for task_id={task_id}",
                        response_json={"task_id": task_id},
                    )
                )
            version = float(max(metrics.versions))

        return self._client.request_json(
            "POST",
            "/recalculate-evaluation",
            json_body={
                "client_id": self._client_id,
                "llm_usecase_id": task_id,
                "version": float(version),
            },
        )  # type: ignore[return-value]

    def get_benchmark_predictions(
        self,
        *,
        task_id: str,
        version: float,
        compare_with_version: float | None = None,
        compare_with_file: bool = False,
        sub_task_type: str | None = None,
        page_size: int = 200,
    ):
        """
        Iterate benchmark model responses via paginated `/model-responses`.

        Yields `ModelResponse` objects until the server returns no more rows.
        """
        page = 1
        while True:
            payload: Dict[str, Any] = {
                "client_id": self._client_id,
                "llm_usecase_id": task_id,
                "version": float(version),
                "page": int(page),
                "page_size": int(page_size),
                "filters": {},
                "comparison_version": compare_with_version,
                "is_version_comparison": bool(compare_with_version is not None),
                "is_file_comparison": bool(compare_with_file),
                "sub_task_type": sub_task_type,
                "primary_model": "student",
            }
            data = self._client.request_json(
                "POST", "/model-responses", json_body=payload
            )
            if not isinstance(data, dict) or not data.get("success"):
                break
            rows = data.get("data") or []
            if not rows:
                break
            for row in rows:
                if isinstance(row, dict):
                    yield ModelResponse.parse_obj(row)
            page += 1

    def request_download(
        self,
        *,
        task_id: str,
        deployed_version: float,
        client_uuid: str,
    ) -> Dict[str, Any]:
        """
        Request model packaging for download.

        This triggers an async process; the backend later emits a `downloadCompletionUpdate`
        event over the `/updates-ws` WebSocket containing `downloadPath` and `filename`.

        Args:
            task_id: Task id.
            deployed_version: Version to download.
            client_uuid: Correlation id used by the download completion event.

        Returns:
            dict: API response (includes `progressInfo` on success).
        """
        payload = {
            "client_id": self._client_id,
            "llm_usecase_id": task_id,
            "deployed_version": deployed_version,
            "client_uuid": client_uuid,
        }
        return self._client.request_json("POST", "/download-model-request", json_body=payload)  # type: ignore[return-value]

    def download_model(
        self,
        *,
        task_id: str,
        version: float,
        output_path: Union[str, Path],
        timeout_s: float = 900.0,
        wait_poll_timeout_s: float = 60.0,
        chunk_size: int = 1024 * 1024,
    ) -> Path:
        """
        Download a packaged model artifact zip to a local path (single-step helper).

        This method handles the full flow:
        1) request packaging (`POST /download-model-request`)
        2) wait for completion (`POST /wait-download-completion`)
        3) stream the zip (`POST /download-model`) into `output_path`

        Args:
            task_id: Task id.
            version: Model version to package and download.
            output_path: Local file path (recommended, ending with `.zip`) OR a directory path.
                - If a directory is provided, the SDK will save `<filename>.zip` inside it.
            timeout_s: Overall timeout for waiting on packaging completion.
            wait_poll_timeout_s: Per-request long-poll timeout to `/wait-download-completion`.
            chunk_size: Streaming chunk size.

        Returns:
            Path: The local file path written.

        Raises:
            TimeoutError: If packaging does not complete within `timeout_s`.
            SpecificAIAPIError: If the API returns an error.
        """
        client_uuid = str(uuid.uuid4())
        req = self.request_download(
            task_id=task_id, deployed_version=version, client_uuid=client_uuid
        )
        # Backend historically returned 200 with {"success": false, ...}; fail fast for SDK users.
        if isinstance(req, dict) and req.get("success") is False:
            msg = str(req.get("message") or "Download request failed")
            # Most common case: wrong task_id or version.
            if "not found" in msg.lower() or "no training" in msg.lower():
                raise NotFoundError(
                    APIErrorDetails(status_code=404, message=msg, response_json=req)
                )
            raise RuntimeError(msg)

        deadline = time.monotonic() + float(timeout_s)
        completion: Optional[Dict[str, Any]] = None

        while time.monotonic() < deadline:
            remaining = deadline - time.monotonic()
            poll_timeout = min(float(wait_poll_timeout_s), max(1.0, remaining))
            try:
                resp = self._client.request_json(
                    "POST",
                    "/wait-download-completion",
                    json_body={"client_uuid": client_uuid, "timeout_s": poll_timeout},
                    timeout_s=poll_timeout + 10.0,
                )
            except NetworkError:
                # Transient network/timeout while long-polling; keep waiting until overall deadline.
                continue

            # On timeout, backend returns a small error payload; keep waiting.
            if (
                isinstance(resp, dict)
                and resp.get("action") != "downloadCompletionUpdate"
            ):
                continue

            # We may receive progress updates without downloadPath/filename. Wait for terminal state.
            if isinstance(resp, dict):
                if resp.get("isError"):
                    completion = resp
                    break
                if resp.get("isCompleted"):
                    completion = resp
                    break
                continue

        if not completion:
            raise TimeoutError(
                f"Timed out waiting for model packaging (task_id={task_id}, version={version})"
            )

        if completion.get("isError"):
            raise RuntimeError(
                completion.get("errorMessage")
                or "Model packaging failed (downloadCompletionUpdate isError=True)"
            )

        download_path = completion.get("downloadPath")
        filename = completion.get("filename")
        resolved_version = completion.get("version", version)
        if not download_path or not filename:
            raise RuntimeError(
                f"Download completion missing downloadPath/filename: {completion}"
            )

        out = Path(output_path)
        # If output_path is a directory, place `<filename>.zip` inside it.
        if out.exists() and out.is_dir():
            out = out / (
                f"{filename}.zip"
                if not str(filename).endswith(".zip")
                else str(filename)
            )
        elif out.suffix == "":
            # Treat a suffix-less path as a directory-like target.
            out.mkdir(parents=True, exist_ok=True)
            out = out / (
                f"{filename}.zip"
                if not str(filename).endswith(".zip")
                else str(filename)
            )
        out.parent.mkdir(parents=True, exist_ok=True)

        stream_resp = self._client.request(
            "POST",
            "/download-model",
            json_body={
                "client_id": self._client_id,
                "llm_usecase_id": task_id,
                "version": float(resolved_version),
                "download_path": download_path,
                "filename": filename,
            },
            stream=True,
            timeout_s=300.0,
        )

        with out.open("wb") as f:
            for chunk in stream_resp.iter_content(chunk_size=chunk_size):
                if chunk:
                    f.write(chunk)
        return out

    def download(
        self,
        *,
        task_id: str,
        version: float,
        download_path: str,
        filename: str,
        output_path: str | Path,
        chunk_size: int = 1024 * 1024,
    ) -> Path:
        """
        Download a packaged model artifact zip.

        Note:
            This endpoint requires `download_path` and `filename` which are typically provided
            by a `downloadCompletionUpdate` event from the `/updates-ws` WebSocket.

        Args:
            task_id: Task id.
            version: Model version.
            download_path: Remote storage path provided by the completion event.
            filename: Zip filename provided by the completion event.
            output_path: Local path to write the zip to.
            chunk_size: Streaming chunk size.

        Returns:
            Path: The local file path written.

        Raises:
            SpecificAIAPIError: If the API returns an error.
        """
        out = Path(output_path)
        out.parent.mkdir(parents=True, exist_ok=True)

        resp = self._client.request(
            "POST",
            "/download-model",
            json_body={
                "client_id": self._client_id,
                "llm_usecase_id": task_id,
                "version": version,
                "download_path": download_path,
                "filename": filename,
            },
            stream=True,
            timeout_s=300.0,
        )

        with out.open("wb") as f:
            for chunk in resp.iter_content(chunk_size=chunk_size):
                if chunk:
                    f.write(chunk)
        return out
