from __future__ import annotations

import json
import base64
from pathlib import Path
from typing import Any, Dict, Optional

from specific_ai.platform.base_client import BaseClient
from specific_ai.platform.schemas import (
    DatasetConfig,
    FileColumnsResponse,
    FileLabelsResponse,
    HuggingFaceDatasetColumnsResponse,
    HuggingFaceDatasetLabelsResponse,
    HuggingFaceDatasetSplitsResponse,
    UploadDatasetResponse,
    UploadHuggingFaceDatasetResponse,
    UploadStatus,
)


DEFAULT_CONFIG = DatasetConfig()


class AssetsManager:
    """Upload datasets / benchmarks to the SpecificAI platform."""

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

    def upload_dataset(
        self,
        *,
        file_path: str | Path,
        task_type: str,
        is_benchmark: bool,
        task_id: Optional[str] = None,
        data_creation_type: Optional[str] = None,
        example_column_name: str,
        label_column_name: str,
        is_manual_model_predictions: bool = False,
        config: DatasetConfig | None = DEFAULT_CONFIG,
    ) -> UploadDatasetResponse:
        """
        Upload a dataset/benchmark file (multipart/form-data).

        Args:
            file_path: Path to the local file to upload.
            task_type: Task type string (e.g. "ClassificationResponse", "NERResponse", ...).
            is_benchmark: Whether this file is a benchmark dataset.
            task_id: Optional task id to associate the dataset with.
            data_creation_type: Optional data creation type value (backend accepts strings).
            example_column_name: Column containing examples.
            label_column_name: Column containing labels.
            config: DatasetConfig object containing the configuration for the dataset.

        Returns:
            UploadDatasetResponse: Upload status and associated dataset lists.

        Raises:
            FileNotFoundError: If `file_path` does not exist.
            SpecificAIAPIError: If the API returns an error.
        """
        path = Path(file_path)
        if not path.exists():
            raise FileNotFoundError(str(path))

        data = {
            "client_id": self._client_id,
            "is_benchmark": str(bool(is_benchmark)).lower(),
            "is_manual_model_predictions": str(
                bool(is_manual_model_predictions)
            ).lower(),
            "task_type": task_type,
            "usecase_id": task_id or "",
            "data_creation_type": data_creation_type or "",
            "label_mappings": json.dumps(config.label_mappings or {}),
            "example_column_name": example_column_name,
            "label_column_name": label_column_name,
            "relations_column_name": config.relations_column_name or "",
            "relations_mappings": json.dumps(config.relations_mappings or {}),
            "classification_labels_delimiter": config.classification_labels_delimiter
            or "",
            "examples_are_with_prompt": str(
                bool(config.examples_are_with_prompt)
            ).lower(),
            "prompt_template": config.prompt_template or "",
            "example_dynamic_field_name": config.example_dynamic_field_name or "",
        }

        # FastAPI collects multiple values with same key into a list; requests supports this via tuples.
        multipart_data: list[tuple[str, str]] = [
            (k, v) for k, v in data.items() if v is not None
        ]
        if config.dynamic_fields_in_prompt_template:
            for v in config.dynamic_fields_in_prompt_template:
                multipart_data.append(("dynamic_fields_in_prompt_template", v))

        with path.open("rb") as f:
            suffix = path.suffix.lower()
            if suffix == ".json":
                files = {"file": (path.name, f, "application/json")}
            elif suffix in (".jsonl", ".ndjson"):
                files = {"file": (path.name, f, "application/x-ndjson")}
            else:
                files = {"file": (path.name, f)}
            resp = self._client.request_json(
                "POST",
                "/upload-dataset",
                data=multipart_data,
                files=files,
            )
        return UploadDatasetResponse.parse_obj(resp)

    def upload_huggingface_dataset(
        self,
        *,
        dataset_id: str,
        label_column: str,
        text_column: str,
        task_type: str,
        train_split: str | None = None,
        test_split: str | None = None,
        task_id: str | None = None,
        data_creation_type: str = "huggingface",
        label_mappings: Optional[Dict[str, Any]] = None,
        is_multilabel: bool = False,
        delimiter: Optional[str] = None,
        dataset_name: str | None = None,
    ) -> UploadHuggingFaceDatasetResponse:
        """
        Upload a Hugging Face dataset and process it in the background.

        This endpoint is asynchronous. It returns a `status_id` that can be polled using
        `get_upload_status(status_id=...)`.

        Notes:
            - The backend expects the user's Hugging Face token to be configured in the platform.
            - The backend can associate the uploaded splits with a task when `task_id` is provided.

        Args:
            dataset_id: Hugging Face dataset identifier (e.g. `"allenai/mini-trec"`).
            label_column: Column name containing labels in the dataset.
            text_column: Column name containing the input text in the dataset.
            task_type: Task type string (e.g. `"ClassificationResponse"`).
            train_split: Split to upload as training dataset (e.g. `"train"`).
            test_split: Split to upload as benchmark dataset (e.g. `"test"`).
            task_id: Optional task id to associate the dataset/benchmark with.
            data_creation_type: Data creation type value (default: `"huggingface"`).
            label_mappings: Optional mapping of labels (e.g. class remapping).
            is_multilabel: Whether this is a multi-label classification dataset.
            delimiter: Optional delimiter for multi-label values.
            dataset_name: Optional platform-side dataset key (Mongo collections,
                dataset_info, ingestion). When omitted, the backend derives the key from
                `dataset_id` (legacy behavior). Use this to import the same HF repo again
                without merging into an existing platform dataset. Note that the backend may
                normalize this value (e.g., character replacement, trimming, and truncation
                to a maximum length), so the stored dataset key may differ from the exact
                string provided here.

        Returns:
            UploadHuggingFaceDatasetResponse: Contains `status_id` for polling.

        Raises:
            ValidationError: If the server rejects the payload (e.g., missing splits).
            SpecificAIAPIError: If the API returns an error.
        """
        payload: Dict[str, Any] = {
            "dataset_id": dataset_id,
            "label_column": label_column,
            "text_column": text_column,
            "label_mappings": label_mappings or {},
            "client_id": self._client_id,
            "is_benchmark": False,
            "task_type": task_type,
            "usecase_id": task_id,
            "data_creation_type": data_creation_type,
            "train_split": train_split,
            "test_split": test_split,
            "is_multilabel": bool(is_multilabel),
            "delimiter": delimiter,
        }
        if dataset_name is not None:
            payload["dataset_name"] = dataset_name

        resp = self._client.request_json(
            "POST",
            "/upload-huggingface-dataset",
            json_body=payload,
        )
        return UploadHuggingFaceDatasetResponse.parse_obj(resp)

    def get_huggingface_dataset_columns(
        self,
        *,
        dataset_id: str,
    ) -> HuggingFaceDatasetColumnsResponse:
        """
        Get column names from a Hugging Face dataset (server-side loading).

        Args:
            dataset_id: Hugging Face dataset identifier (e.g. `"allenai/mini-trec"`).

        Returns:
            HuggingFaceDatasetColumnsResponse: `columns` on success, otherwise `error`.
        """
        resp = self._client.request_json(
            "POST",
            "/get-hf-dataset-columns",
            json_body={"dataset_id": dataset_id, "client_id": self._client_id},
        )
        return HuggingFaceDatasetColumnsResponse.parse_obj(resp)

    def get_huggingface_dataset_labels(
        self,
        *,
        dataset_id: str,
        label_column: str,
        is_multilabel: bool = False,
        delimiter: Optional[str] = None,
    ) -> HuggingFaceDatasetLabelsResponse:
        """
        Get unique labels from a Hugging Face dataset label column (server-side loading).

        Args:
            dataset_id: Hugging Face dataset identifier (e.g. `"allenai/mini-trec"`).
            label_column: Column name containing labels.
            is_multilabel: Whether the label column contains multi-label values.
            delimiter: Optional delimiter for multi-label values.

        Returns:
            HuggingFaceDatasetLabelsResponse: `labels` on success, otherwise `error`.
        """
        resp = self._client.request_json(
            "POST",
            "/get-hf-dataset-labels",
            json_body={
                "dataset_id": dataset_id,
                "label_column": label_column,
                "client_id": self._client_id,
                "is_multilabel": bool(is_multilabel),
                "delimiter": delimiter,
            },
        )
        return HuggingFaceDatasetLabelsResponse.parse_obj(resp)

    def get_huggingface_dataset_splits(
        self,
        *,
        dataset_id: str,
    ) -> HuggingFaceDatasetSplitsResponse:
        """
        Get available split names for a Hugging Face dataset (server-side loading).

        Args:
            dataset_id: Hugging Face dataset identifier (e.g. `"CogComp/trec"`).

        Returns:
            HuggingFaceDatasetSplitsResponse: `splits` on success, otherwise `error`.
        """
        resp = self._client.request_json(
            "POST",
            "/get-hf-dataset-splits",
            json_body={"dataset_id": dataset_id, "client_id": self._client_id},
        )
        return HuggingFaceDatasetSplitsResponse.parse_obj(resp)

    def delete_dataset(
        self,
        *,
        dataset: str,
        is_benchmark: bool,
        is_manual_model_predictions: bool = False,
    ) -> dict:
        """
        Delete a dataset/benchmark (drops metadata + underlying collection).

        Args:
            dataset: Dataset name (note: backend stores dataset name without file extension).
            is_benchmark: Whether this dataset is a benchmark.
            is_manual_model_predictions: Whether this dataset is an external predictions file.

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

        Raises:
            SpecificAIAPIError: If the API returns an error.
        """
        return self._client.request_json(
            "DELETE",
            "/datasets",
            json_body={
                "client_id": self._client_id,
                "dataset": dataset,
                "is_benchmark": is_benchmark,
                "is_manual_model_predictions": is_manual_model_predictions,
            },
        )  # type: ignore[return-value]

    def get_low_confidence_samples(
        self,
        *,
        dataset: str,
        label: str,
        task_id: str,
        is_benchmark: bool,
        k: int = 30,
    ) -> Dict[str, Any]:
        """
        Fetch "improvement bucket" samples for re-annotation (Active Learning).

        This returns up to `k` **unvalidated** records from the given dataset and label bucket.
        On the backend, results are prioritized by a rank score and may fall back to semantic
        similarity when there aren't enough samples in the bucket.

        Args:
            dataset: Dataset name (as returned by upload endpoints / shown in the UI).
            label: The label bucket to fetch samples for.
            task_id: Task id (llm_usecase_id).
            is_benchmark: Whether the dataset is a benchmark dataset.
            k: Max number of samples to return (default: 30).

        Returns:
            dict: API response with shape:
                {
                  "success": bool,
                  "examples": [
                    {
                      "id": str,
                      "record_id": str,
                      "llm_usecase_id": str,
                      "request_example": Any,
                      "labels": list[Any],
                      "is_validated": bool,
                      "response": dict[str, Any],
                    },
                    ...
                  ]
                }

        Raises:
            SpecificAIAPIError: If the API returns an error.
        """
        payload = {
            "dataset": dataset,
            "label": label,
            "client_id": self._client_id,
            "is_benchmark": bool(is_benchmark),
            "k": int(k),
            "llm_usecase_id": task_id,
        }
        return self._client.request_json(
            "POST",
            "/get-top-ranked-examples",
            json_body=payload,
        )  # type: ignore[return-value]

    def get_upload_status(self, *, status_id: str) -> UploadStatus:
        """
        Get dataset upload processing status.

        Args:
            status_id: The `status_id` returned by `upload_dataset()` / `upload_huggingface_dataset()`.

        Returns:
            UploadStatus: Current progress/status.

        Raises:
            SpecificAIAPIError: If the API returns an error (e.g. 404 if not found).
        """
        data = self._client.request_json("GET", f"/upload-status/{status_id}")
        return UploadStatus.parse_obj(data)

    def get_file_columns(
        self,
        *,
        file_path: str | Path,
        file_type: Optional[str] = None,
    ) -> FileColumnsResponse:
        """
        Infer column names from a local file using the backend's parsing logic.

        This is useful to validate your CSV/JSON schema before uploading, and to choose
        the right `example_column_name` / `label_column_name`.

        Args:
            file_path: Local path to a CSV or JSON file.
            file_type: Optional override ("csv" or "json"). If not provided, inferred from extension.

        Returns:
            FileColumnsResponse
        """
        path = Path(file_path)
        content = path.read_bytes()
        inferred = (path.suffix or "").lstrip(".").lower()
        ft = (file_type or inferred or "").lower()
        payload = {
            "filename": path.name,
            "client_id": self._client_id,
            "file_content": base64.b64encode(content).decode("utf-8"),
            "file_type": ft,
        }
        data = self._client.request_json("POST", "/get-file-columns", json_body=payload)
        return FileColumnsResponse.parse_obj(data)

    def get_file_labels(
        self,
        *,
        file_path: str | Path,
        label_column: str,
        task_type: str,
        delimiter: Optional[str] = None,
        file_type: Optional[str] = None,
    ) -> FileLabelsResponse:
        """
        Infer unique labels from a file column using the backend's parsing logic.

        Args:
            file_path: Local path to a CSV or JSON file.
            label_column: Column name containing labels.
            task_type: Task type string (e.g. "ClassificationResponse", "NERResponse", ...).
            delimiter: Optional delimiter for multilabel CSV cells.
            file_type: Optional override ("csv" or "json"). If not provided, inferred from extension.

        Returns:
            FileLabelsResponse
        """
        path = Path(file_path)
        content = path.read_bytes()
        inferred = (path.suffix or "").lstrip(".").lower()
        ft = (file_type or inferred or "").lower()
        payload = {
            "filename": path.name,
            "client_id": self._client_id,
            "label_column": label_column,
            "file_content": base64.b64encode(content).decode("utf-8"),
            "file_type": ft,
            "task_type": task_type,
            "delimiter": delimiter,
        }
        data = self._client.request_json("POST", "/get-file-labels", json_body=payload)
        return FileLabelsResponse.parse_obj(data)
