from __future__ import annotations

from typing import Any, List, Literal, Optional

from specific_ai.platform.schemas import ModelProvider

from specific_ai.platform.base_client import BaseClient
from specific_ai.platform.schemas import (
    CreateTasksResponse,
    Task,
    TaskCreate,
    TaskGroup,
)


class TaskManager:
    """
    Manage Tasks.

    """

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

    def create(
        self,
        *,
        project_name: str,
        tasks: List[TaskCreate],
        project_goal: Optional[str] = None,
    ) -> CreateTasksResponse:
        """
        Create one or more tasks under a project.

        Args:
            project_name: Project name.
            tasks: List of task definitions to create.
            project_goal: Optional project goal/description.

        Returns:
            CreateTasksResponse: Contains `success` and `created_task_ids`.

        Raises:
            SpecificAIAPIError: If the API returns an error.
        """
        payload = {
            "client_id": self._client_id,
            "group_name": project_name,
            "project_goal": project_goal,
            "tasks": [t.dict(by_alias=True) for t in tasks],
        }
        data = self._client.request_json(
            "POST", "/create_llm_usecase", json_body=payload
        )
        return CreateTasksResponse.parse_obj(data)

    def list(self) -> List[TaskGroup]:
        """
        List tasks grouped by project and subgroup.

        Args:
        Returns:
            List[TaskGroup]: Groups containing sub-groups and tasks.

        Raises:
            SpecificAIAPIError: If the API returns an error.
        """
        data = self._client.request_json(
            "POST", "/llm_usecases", json_body={"value": self._client_id}
        )
        # Endpoint returns a JSON list.
        return [TaskGroup.parse_obj(item) for item in data]  # type: ignore[arg-type]

    def get(self, *, task_id: str) -> Task:
        """
        Fetch a single task by id.

        Args:
            task_id: Task id.

        Returns:
            Task: The task record.

        Raises:
            SpecificAIAPIError: If the API returns an error.
        """
        data = self._client.request_json(
            "POST",
            "/llm_usecase_info",
            json_body={"client_id": self._client_id, "llm_usecase_id": task_id},
        )
        # Shape: {"llm_usecase": {...}} or {"error": "..."}
        if isinstance(data, dict) and "llm_usecase" in data:
            return Task.parse_obj(data["llm_usecase"])
        return Task.parse_obj(data)

    def edit(
        self,
        *,
        task_id: str,
        task_name: Optional[str] = None,
        task_goal: Optional[str] = None,
        task_type: Optional[str] = None,
        project_name: Optional[str] = None,
        is_multilabel: Optional[bool] = None,
        task_languages: Optional[List[str]] = None,
    ) -> dict:
        """
        Edit mutable task fields.

        Note:
            Tasks are largely managed by the platform (datasets, benchmarks, status,
            versions, etc.). This method only updates the fields supported by the
            backend `POST /update_llm_usecase` endpoint.

        Args:
            task_id: Task id.
            task_name: Optional new task name. Defaults to current value.
            task_goal: Optional new goal/description. Defaults to current value.
            task_type: Optional task type string. Defaults to current value.
            project_name: Optional project name. Defaults to current value.
            is_multilabel: Optional multilabel flag. Defaults to current value.
            task_languages: Optional list of language codes. Defaults to current value.

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

        Raises:
            SpecificAIAPIError: If the API returns an error.
            ValueError: If required fields are missing.
        """
        current = self.get(task_id=task_id)
        if not (task_name or current.task_name):
            raise ValueError("task_name is required")
        if not (task_type or current.task_type):
            raise ValueError("task_type is required")

        payload = {
            "usecase_id": task_id,
            "client_id": self._client_id,
            "usecase_name": task_name or current.task_name,
            "usecase_goal": task_goal if task_goal is not None else current.task_goal,
            "task_type": task_type or current.task_type,
            "group_name": (
                project_name if project_name is not None else current.project_name
            ),
            "is_multilabel": (
                bool(is_multilabel)
                if is_multilabel is not None
                else bool(current.is_multilabel)
            ),
            "classify_speakers": False,
            "task_languages": (
                task_languages
                if task_languages is not None
                else list(current.task_languages or [])
            ),
        }
        return self._client.request_json(
            "POST", "/update_llm_usecase", json_body=payload
        )  # type: ignore[return-value]

    def save_teacher_and_prompt(
        self,
        *,
        task_id: str,
        teacher: str,
        prompt: str,
        provider: Optional[ModelProvider] = ModelProvider.openai,
        teacher_display: Optional[str] = None,
    ) -> dict[str, Any]:
        """
        Save the teacher model and prompt for a task.

        This is required before starting training. The backend also schedules
        background example extraction based on the prompt.

        Args:
            task_id: Task id.
            teacher: Commercial/teacher model name (e.g. "gpt-4o-mini").
            prompt: Prompt/instruction template for the teacher.
            provider: Optional provider identifier (backend accepts string).
            teacher_display: Optional display name for UI.

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

        Raises:
            SpecificAIAPIError: If the API returns an error.
        """
        payload = {
            "client_id": self._client_id,
            "llm_usecase_id": task_id,
            "teacher": teacher,
            "prompt": prompt,
            "provider": (
                provider.value if isinstance(provider, ModelProvider) else provider
            ),
            "teacher_display": teacher_display,
        }
        return self._client.request_json(
            "POST", "/save_teacher_and_prompt", json_body=payload
        )  # type: ignore[return-value]

    def set_comparison_config(
        self,
        *,
        task_id: str,
        compare_with_teacher: bool = False,
        compare_with_results_file: bool = False,
        response_parser: Optional[str] = None,
    ) -> dict[str, Any]:
        """
        Configure comparison mode for a task (doc-friendly helper).

        This is a high-level wrapper that:
        - Optionally validates & saves a `response_parser` function (required for teacher comparisons)
        - Sets the backend comparison mode option

        Mapping:
            - teacher only -> option="teacher"
            - results file only -> option="manual_model_predictions"
            - both -> option="all"
            - neither -> option="none"
        """
        if compare_with_teacher and not response_parser:
            raise ValueError(
                "`response_parser` is required when compare_with_teacher=True"
            )

        # Persist parser function (backend endpoint is named `validate-code`).
        if compare_with_teacher and response_parser:
            task = self.get(task_id=task_id)
            self._client.request_json(
                "POST",
                "/validate-code",
                json_body={
                    "code": response_parser,
                    "raw_response": "example",
                    "client_id": self._client_id,
                    "llm_usecase_id": task_id,
                    "task_type": task.task_type,
                },
            )

        if compare_with_teacher and compare_with_results_file:
            option = "all"
        elif compare_with_teacher:
            option = "teacher"
        elif compare_with_results_file:
            option = "manual_model_predictions"
        else:
            option = "none"

        return self.update_comparison_mode_option(
            task_id=task_id, comparison_mode_option=option  # type: ignore[arg-type]
        )

    def update_comparison_mode_option(
        self,
        *,
        task_id: str,
        comparison_mode_option: Literal[
            "none", "teacher", "manual_model_predictions", "all"
        ],
    ) -> dict[str, Any]:
        """
        Update comparison mode option for a task.

        Comparison mode controls which prediction sources are used for evaluation comparisons.

        Allowed options:
            - "none"
            - "teacher"
            - "manual_model_predictions"
            - "all"

        Args:
            task_id: Task id.
            comparison_mode_option: One of: "none", "teacher", "manual_model_predictions", "all".

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

        Raises:
            SpecificAIAPIError: If the API returns an error.
        """
        # The backend also requires the boolean comparison_mode flag to be enabled
        # when the option is not "none".
        self.update_comparison_mode(
            task_id=task_id, comparison_mode=(comparison_mode_option != "none")
        )

        payload = {
            "client_id": self._client_id,
            "llm_usecase_id": task_id,
            "comparison_mode_option": comparison_mode_option,
        }
        return self._client.request_json(
            "POST", "/update_comparison_mode_option", json_body=payload
        )  # type: ignore[return-value]

    def update_comparison_mode(
        self,
        *,
        task_id: str,
        comparison_mode: bool,
    ) -> dict[str, Any]:
        """
        Update comparison mode boolean flag for a task.

        Args:
            task_id: Task id.
            comparison_mode: Enable/disable comparison mode.

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

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

    def delete(self, *, task_id: str) -> dict:
        """
        Delete a task.

        Args:
            task_id: Task id.

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

        Raises:
            SpecificAIAPIError: If the API returns an error.
        """
        return self._client.request_json(
            "POST",
            "/delete_llm_usecase",
            json_body={"client_id": self._client_id, "usecase_id": task_id},
        )  # type: ignore[return-value]
