from __future__ import annotations

from typing import Annotated, List, Literal, Optional, Union
from uuid import UUID

from pydantic import BaseModel, Discriminator, Field

from schemas.base_experiment_schemas import (
    BaseConfigResult,
    BaseCreateExperimentRequest,
    BaseExperimentDetail,
    BaseExperimentSummary,
    BaseResult,
    BaseTestCase,
    DatasetColumnVariableSource,
    EvalResultSummary,
)
from schemas.common_schemas import BasePaginationResponse
from schemas.request_schemas import (
    RagSearchSettingConfigurationRequestTypes,
)
from schemas.response_schemas import (
    RagProviderQueryResponse,
    RagSearchSettingConfigurationResponseTypes,
)


class SavedRagConfig(BaseModel):
    """Configuration for a saved RAG setting configuration version"""

    type: Literal["saved"] = "saved"
    setting_configuration_id: UUID = Field(
        description="ID of the RAG search setting configuration",
    )
    version: int = Field(description="Version of the RAG search setting configuration")
    query_column: DatasetColumnVariableSource = Field(
        description="Dataset column to use as the RAG search query",
    )


class UnsavedRagConfig(BaseModel):
    """Configuration for an unsaved RAG search setting"""

    type: Literal["unsaved"] = "unsaved"
    unsaved_id: Optional[UUID] = Field(
        default=None,
        description="Unique identifier for this unsaved RAG configuration (generated by backend if not provided)",
    )
    rag_provider_id: UUID = Field(
        description="ID of the RAG provider to use for this search",
    )
    settings: RagSearchSettingConfigurationRequestTypes = Field(
        description="RAG search settings configuration",
    )
    query_column: DatasetColumnVariableSource = Field(
        description="Dataset column to use as the RAG search query",
    )


RagConfig = Annotated[
    Union[SavedRagConfig, UnsavedRagConfig],
    Discriminator("type"),
]


class UnsavedRagConfigResponse(BaseModel):
    """Configuration for an unsaved RAG search setting (response version)"""

    type: Literal["unsaved"] = "unsaved"
    unsaved_id: Optional[UUID] = Field(
        default=None,
        description="Unique identifier for this unsaved RAG configuration (generated by backend if not provided)",
    )
    rag_provider_id: UUID = Field(
        description="ID of the RAG provider to use for this search",
    )
    settings: RagSearchSettingConfigurationResponseTypes = Field(
        description="RAG search settings configuration",
    )
    query_column: DatasetColumnVariableSource = Field(
        description="Dataset column to use as the RAG search query",
    )


RagConfigResponse = Annotated[
    Union[SavedRagConfig, UnsavedRagConfigResponse],
    Discriminator("type"),
]


# RAG Experiment schemas
class RagExperimentSummary(BaseExperimentSummary):
    """Summary of a RAG experiment"""

    rag_configs: List[RagConfigResponse] = Field(
        description="List of RAG configurations being tested",
    )


class CreateRagExperimentRequest(BaseCreateExperimentRequest):
    """Request to create a new RAG experiment"""

    rag_configs: List[RagConfig] = Field(
        description="List of RAG configurations to test. Each config specifies which dataset column to use as the query.",
    )


class RagEvalResultSummaries(BaseModel):
    """Summary of evaluation results for a RAG configuration"""

    rag_config_key: str | None = Field(
        default=None,
        description="RAG config key: 'saved:setting_config_id:version' for saved, 'unsaved:uuid' for unsaved",
    )
    rag_config_type: str | None = Field(
        default=None,
        description="Type: 'saved' or 'unsaved'",
    )
    setting_configuration_id: UUID | None = Field(
        default=None,
        description="ID of the RAG search setting configuration (for saved configs only)",
    )
    setting_configuration_version: int | None = Field(
        default=None,
        description="Version of the RAG search setting configuration (for saved configs only)",
    )
    eval_results: list[EvalResultSummary] = Field(
        description="Results for each evaluation run on this RAG configuration",
    )


class RagSummaryResults(BaseModel):
    """Summary results across all RAG configurations and evaluations"""

    rag_eval_summaries: list[RagEvalResultSummaries] = Field(
        description="Summary for each RAG configuration tested",
    )


class RagExperimentDetail(BaseExperimentDetail):
    """Detailed information about a RAG experiment"""

    rag_configs: List[RagConfigResponse] = Field(
        description="List of RAG configurations being tested",
    )
    summary_results: RagSummaryResults = Field(
        description="Summary of results across all test cases",
    )


# Pagination schemas
class RagExperimentListResponse(BasePaginationResponse):
    """Paginated list of RAG experiments"""

    data: list[RagExperimentSummary] = Field(
        description="List of RAG experiment summaries",
    )


# Test case / result schemas
class RagSearchOutput(BaseModel):
    """Output from a RAG search execution"""

    response: RagProviderQueryResponse = Field(description="RAG search response")


class RagResult(BaseResult):
    """Results from a RAG search execution with evals"""

    rag_config_key: str = Field(
        description="RAG config key: 'saved:setting_config_id:version' for saved, 'unsaved:uuid' for unsaved",
    )
    rag_config_type: str = Field(
        description="Type: 'saved' or 'unsaved'",
    )
    setting_configuration_id: UUID | None = Field(
        default=None,
        description="ID of the RAG search setting configuration (for saved configs only)",
    )
    setting_configuration_version: int | None = Field(
        default=None,
        description="Version of the RAG search setting configuration (for saved configs only)",
    )
    query_text: str = Field(description="Query text used for the search")
    output: Optional[RagSearchOutput] = Field(
        default=None,
        description="Output from the RAG search (None if not yet executed)",
    )


class RagTestCase(BaseTestCase):
    """Individual test case result for RAG experiment"""

    rag_results: list[RagResult] = Field(
        description="Results for each RAG configuration tested",
    )


class RagTestCaseListResponse(BasePaginationResponse):
    """Paginated list of RAG test cases"""

    data: list[RagTestCase] = Field(description="List of test cases")


class RagConfigResult(BaseConfigResult):
    """Result for a specific RAG configuration within a test case"""

    query_text: str = Field(description="Query text used for the search")
    output: Optional[RagSearchOutput] = Field(
        default=None,
        description="Output from the RAG search (None if not yet executed)",
    )


class RagConfigResultListResponse(BasePaginationResponse):
    """Paginated list of results for a specific RAG configuration"""

    data: list[RagConfigResult] = Field(
        description="List of results for the RAG configuration",
    )
