from __future__ import annotations

from enum import Enum
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field


class SDKBaseModel(BaseModel):
    """Base pydantic model for SDK responses (keeps forward-compat via extra fields)."""

    class Config:
        extra = "allow"
        populate_by_name = True


class ModelProvider(str, Enum):
    openai = "openai"
    bedrock_anthropic = "bedrock-anthropic"
    local = "local"


class Task(SDKBaseModel):
    """A Task in the SpecificAI platform."""

    task_id: Optional[str] = Field(default=None, alias="_id")
    client_id: Optional[str] = None
    project_name: Optional[str] = Field(default=None, alias="group_name")
    sub_group_name: Optional[str] = None

    task_name: Optional[str] = Field(default=None, alias="usecase_name")
    task_goal: Optional[str] = Field(default=None, alias="usecase_goal")
    task_type: Optional[str] = None

    datasets: List[str] = Field(default_factory=list)
    benchmarks: List[str] = Field(default_factory=list)

    prompt: Optional[str] = None
    teacher_model: Optional[str] = None
    parser_function: Optional[str] = None

    comparison_mode: Optional[bool] = None
    comparison_mode_option: Optional[str] = None

    is_multilabel: Optional[bool] = None
    task_languages: List[str] = Field(default_factory=list)

    version: Optional[float] = None
    status: Optional[str] = None


class TaskCreate(SDKBaseModel):
    """Payload for creating a task inside a project."""

    task_name: str = Field(alias="usecase_name")
    task_type: str
    task_goal: Optional[str] = Field(default=None, alias="usecase_goal")

    comparison_mode: bool = False
    datasets: List[str] = Field(default_factory=list)
    benchmarks: List[str] = Field(default_factory=list)

    prompt: Optional[str] = None
    teacher_model: Optional[str] = None
    parser_function: Optional[str] = None

    is_multilabel: bool = False
    task_languages: List[str] = Field(default_factory=list)

    manual_model_predictions: List[str] = Field(default_factory=list)


class TaskSubGroup(SDKBaseModel):
    sub_group_name: str = ""
    tasks: List[Task] = Field(default_factory=list, alias="usecases")


class TaskGroup(SDKBaseModel):
    project_name: str = Field(alias="group_name")
    project_goal: Optional[str] = None
    sub_groups: List[TaskSubGroup] = Field(default_factory=list)

    def iter_tasks(self) -> List[Task]:
        """Flatten all tasks across sub-groups."""
        out: List[Task] = []
        for sg in self.sub_groups:
            out.extend(sg.tasks or [])
        return out


class TrainingJob(SDKBaseModel):
    """A training job."""

    id: Optional[str] = None
    task_id: Optional[str] = Field(default=None, alias="llm_usecase_id")
    distillation_event_id: Optional[str] = None
    status: Optional[str] = None
    progress: Optional[int] = None
    created_datetime: Optional[str] = None
    finished_datetime: Optional[str] = None
    version: Optional[float] = None
    created_by_user: Optional[dict] = None


class ModelMetrics(SDKBaseModel):
    """Evaluation metrics returned by `POST /evaluation`."""

    success: Optional[bool] = None
    versions: List[float] = Field(default_factory=list)
    all_metrics: Dict[str, Any] = Field(default_factory=dict)
    comparativeMetrics: Dict[str, Any] = Field(default_factory=dict)
    confusionMatrix: Dict[str, Any] = Field(default_factory=dict)


class CreateTasksResponse(SDKBaseModel):
    success: bool
    created_task_ids: List[str] = Field(default_factory=list, alias="created_usecases")


class UploadDatasetResponse(SDKBaseModel):
    status: str
    status_id: str
    current_datasets: List[str] = Field(default_factory=list)
    current_benchmarks: List[str] = Field(default_factory=list)


class DatasetConfig(SDKBaseModel):
    """
    Dataset upload configuration (optional) used by `AssetsManager.upload_dataset(...)`.

    This mirrors the documentation and keeps forward-compatibility via `extra = "allow"`.
    """

    label_mappings: Optional[Dict[str, Any]] = None
    relations_mappings: Optional[Dict[str, Any]] = None
    relations_column_name: Optional[str] = None
    classification_labels_delimiter: Optional[str] = None

    examples_are_with_prompt: bool = False
    prompt_template: Optional[str] = None
    dynamic_fields_in_prompt_template: Optional[List[str]] = None
    example_dynamic_field_name: Optional[str] = None


class UploadHuggingFaceDatasetResponse(SDKBaseModel):
    """Response returned by `POST /upload-huggingface-dataset`."""

    success: bool = True
    status_id: str
    message: Optional[str] = None


class UploadStatus(SDKBaseModel):
    """Status object returned by `GET /upload-status/{status_id}`."""

    filename: str
    total_rows: int
    processed_rows: int
    status: str  # "processing" | "completed" | "failed"
    error: Optional[str] = None
    started_at: Optional[str] = None
    completed_at: Optional[str] = None


class FileColumnsResponse(SDKBaseModel):
    success: bool
    columns: List[str] = Field(default_factory=list)
    error: Optional[str] = None


class FileLabelsResponse(SDKBaseModel):
    success: bool
    labels: List[str] = Field(default_factory=list)
    label_examples: Dict[str, Any] = Field(default_factory=dict)
    error: Optional[str] = None
    preview_truncated: bool = False
    preview_size_bytes: Optional[int] = None
    original_size_bytes: Optional[int] = None
    total_label_count: Optional[int] = None


class StartTrainingResponse(SDKBaseModel):
    distillation_event_id: Optional[str] = None
    training_created: Optional[bool] = None


class HuggingFaceDatasetColumnsResponse(SDKBaseModel):
    """Response returned by `POST /get-hf-dataset-columns`."""

    success: bool
    columns: List[str] = Field(default_factory=list)
    error: Optional[str] = None


class HuggingFaceDatasetLabelsResponse(SDKBaseModel):
    """Response returned by `POST /get-hf-dataset-labels`."""

    success: bool
    labels: List[str] = Field(default_factory=list)
    error: Optional[str] = None


class HuggingFaceDatasetSplitsResponse(SDKBaseModel):
    """Response returned by `POST /get-hf-dataset-splits`."""

    success: bool
    splits: List[str] = Field(default_factory=list)
    error: Optional[str] = None


class ModelResponse(SDKBaseModel):
    """
    One model response row returned by `/model-responses`.

    Field names are kept aligned with API reference docs, while allowing extra backend fields.
    """

    request_example: Dict[str, Any] = Field(default_factory=dict)
    expected_response: Dict[str, Any] = Field(default_factory=dict)
    model_response: Dict[str, Any] = Field(default_factory=dict)
    comparison_model_response: Optional[Dict[str, Any]] = None
    teacher_response: Optional[Dict[str, Any]] = None
