from __future__ import annotations

from typing import Any, Dict, List, Optional, Literal

from specific_ai.platform.base_client import BaseClient
from specific_ai.platform.schemas import StartTrainingResponse, TrainingJob, Task


class TrainingManager:
    """Start and inspect training jobs."""

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

    def start(
        self,
        *,
        task_id: str,
        base_model_name: str = "bert-base-uncased",
        train_ratio: int = 70,
        validation_ratio: int = 15,
        metric_optimization: str = "f_beta_score",
        beta_for_fscore: float = 1.0,
        hardware_device: Literal["cpu", "gpu"] = "gpu",
        training_arguments: Optional[dict] = None,
        commit_hash: Optional[str] = None,
    ) -> StartTrainingResponse:
        """
        Start a training run for a task id.

        This method is intentionally high-level: users only provide a `task_id` and
        optionally a few training knobs. The SDK fetches the task, computes the next
        version, builds the distillation + training payloads, and calls the API.

        This mirrors the web UI flow:
        - POST /create_distillation_event
        - POST /create_training

        Args:
            task_id: Task id to train.
            base_model_name: Base model name (UI default: "bert-base-uncased").
            train_ratio: Train split ratio (0-100).
            validation_ratio: Validation split ratio (0-100).
            metric_optimization: Metric name used for model selection (UI default: "f_beta_score").
            beta_for_fscore: Used when `metric_optimization == "f_beta_score"`.
            hardware_device: "gpu" or "cpu".
            training_arguments: Additional training arguments.
            commit_hash: Optional git commit hash to use as the training image tag (e.g. for job_runner).

        Returns:
            StartTrainingResponse: Includes `distillation_event_id` and `training_created`.

        Raises:
            SpecificAIAPIError: If the API returns an error.
            ValueError: If the task is not ready for training (missing datasets/teacher/prompt).
        """
        task = self._get_task(task_id=task_id)
        self._validate_task_ready_for_training(task)

        new_version = float((task.version or 0) + 1)
        test_ratio = 100 - int(train_ratio) - int(validation_ratio)
        if test_ratio < 0:
            raise ValueError(
                f"Invalid split ratios: train_ratio({train_ratio}) + validation_ratio({validation_ratio}) must be <= 100."
            )

        model_eval_args: Dict[str, Any] = {"optimization_metric": metric_optimization}
        if metric_optimization == "f_beta_score":
            model_eval_args["f_score_metric_beta_coefficient"] = beta_for_fscore

        distillation_event: Dict[str, Any] = {
            "client_id": self._client_id,
            "llm_usecase_id": task_id,
            "datasets": task.datasets,
            "benchmarks": task.benchmarks,
            "teacher_model": task.teacher_model,
            "task_type": task.task_type,
            "base_distilled_model": base_model_name,
            "instruction": task.prompt,
            "version": new_version,
            "parser_function": task.parser_function,
            "hardware_device": hardware_device,
            "comparison_mode": bool(task.comparison_mode),
            "comparison_mode_option": task.comparison_mode_option,
            "is_multilabel": bool(task.is_multilabel),
            "model_evaluation_arguments": model_eval_args,
            "trainer_arguments": training_arguments or {},
        }
        if commit_hash is not None:
            distillation_event["commit_hash"] = commit_hash

        training: Dict[str, Any] = {
            "client_id": self._client_id,
            "version": new_version,
            "llm_usecase_id": task_id,
            "base_model_name": base_model_name,
            "distillation_event_id": "",  # filled after distillation is created
            "training_params": {
                "train_ratio": int(train_ratio),
                "validation_ratio": int(validation_ratio),
                "test_ratio": int(test_ratio),
                "args": {"metric_for_best_model": metric_optimization},
            },
            "hardware_device": hardware_device,
            "model_evaluation_arguments": model_eval_args,
        }

        return self.start_with_payloads(
            distillation_event=distillation_event, training=training
        )

    def stop(self, *, task_id: str) -> Dict[str, Any]:
        """
        Stop a current training run for a task id.

        Args:
            task_id: Task id to stop training for.

        Returns:
            Dict[str, Any]: API response (typically {"success": bool, "message": str}).
        """
        training = self.get_latest(task_id=task_id)
        if not training:
            raise ValueError(f"No training found for task id: {task_id}")
        return self._client.request_json(
            "POST",
            "/delete_training",
            json_body={
                "client_id": self._client_id,
                "distillation_event_id": training.distillation_event_id,
            },
        )

    def start_with_payloads(
        self,
        *,
        distillation_event: Dict[str, Any],
        training: Dict[str, Any],
    ) -> StartTrainingResponse:
        """
        Advanced: start training with explicit payloads.

        Prefer `start(task_id=...)` unless you need full control.

        Args:
            distillation_event: Dict payload compatible with backend DistillationEvent.
            training: Dict payload compatible with backend Training.

        Returns:
            StartTrainingResponse
        """
        dist = self._client.request_json(
            "POST", "/create_distillation_event", json_body=distillation_event
        )
        dist_id = None
        if isinstance(dist, dict):
            dist_id = dist.get("distillation_event_id")

        if dist_id:
            training = dict(training)
            training["distillation_event_id"] = dist_id

        train = self._client.request_json(
            "POST", "/create_training", json_body=training
        )
        created = bool(isinstance(train, dict) and train.get("success") is True)
        return StartTrainingResponse(
            distillation_event_id=dist_id, training_created=created
        )

    def list(self) -> List[TrainingJob]:
        """
        List all trainings for a client.

        Returns:
            List[TrainingJob]: All known training jobs.

        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 []
        return [TrainingJob.parse_obj(t) for t in trainings]

    def get_status(self, *, distillation_event_id: str) -> Optional[TrainingJob]:
        """
        Get status for a specific training, by distillation_event_id.

        Args:
            distillation_event_id: Distillation event id used as training identifier.

        Returns:
            TrainingJob | None: Matching training job if present.

        Raises:
            SpecificAIAPIError: If the API returns an error.
        """
        for t in self.list():
            if t.distillation_event_id == distillation_event_id:
                return t
        return None

    def get(self, *, distillation_event_id: str) -> Optional[TrainingJob]:
        """
        Alias for `get_status`.

        Args:
            distillation_event_id: Distillation event id used as training identifier.

        Returns:
            TrainingJob | None
        """
        return self.get_status(distillation_event_id=distillation_event_id)

    def get_latest(self, *, task_id: str) -> Optional[TrainingJob]:
        """
        Get the latest training for a given task id.

        "Latest" is chosen by highest `version` (fallback: first match).

        Args:
            task_id: Task id.

        Returns:
            TrainingJob | None
        """
        matches = [t for t in self.list() if t.task_id == task_id]
        if not matches:
            return None
        # Prefer highest version when present
        with_version = [t for t in matches if t.version is not None]
        if with_version:
            return max(with_version, key=lambda t: float(t.version))  # type: ignore[arg-type]
        return matches[0]

    def _get_task(self, *, task_id: str) -> Task:
        data = self._client.request_json(
            "POST",
            "/llm_usecase_info",
            json_body={"client_id": self._client_id, "llm_usecase_id": task_id},
        )
        if isinstance(data, dict) and "llm_usecase" in data:
            return Task.parse_obj(data["llm_usecase"])
        return Task.parse_obj(data)

    def _validate_task_ready_for_training(self, task: Task) -> None:
        # Mirrors the UI gating rules (see frontend/src/utils/trainingUtils.js)
        if not (task.datasets and len(task.datasets) > 0):
            raise ValueError(
                "Task is not ready for training: no datasets are associated. Upload a dataset and link it to this task first."
            )
        if not task.prompt or not task.teacher_model:
            raise ValueError(
                "Task is not ready for training: missing `prompt` and/or `teacher_model`. Configure Model Setup in the UI (or via API) first."
            )
        if (
            bool(task.comparison_mode)
            and (task.comparison_mode_option in ("all", "teacher"))
            and not task.parser_function
        ):
            updated_task = self.tasks.get(task_id=task.task_id)
            if not updated_task.parser_function:
                raise ValueError(
                    "Task is not ready for training: comparison_mode requires `parser_function` when comparing with teacher/all."
                )
