from __future__ import annotations

from datetime import datetime
from typing import Optional, Dict, Any
from warnings import warn

from pydantic import ValidationError, Field, validator

from scale_egp.sdk.enums import EvaluationStatus, TestCaseSchemaType
from scale_egp.sdk.types.evaluation_dataset_test_cases import ExtraInfo
from scale_egp.utils.model_utils import Entity, BaseModel


class GenerationTestCaseResultData(BaseModel):
    """
    A data model representing the data of a TestcaseResult with a GENERATION schema type.

    Attributes:
        generation_output: The output of the generation model
        generation_extra_info: The extra info of the generation model
    """

    generation_output: str
    generation_extra_info: Optional[ExtraInfo] = None


class TestCaseResult(Entity):
    """
    A data model representing a test case result.

    Attributes:
        id: The ID of the test case result
        status: The status of the test case result
        application_spec_id: The ID of the application spec that the test case result is for
        evaluation_id: The ID of the evaluation that the test case result is for
        evaluation_dataset_id: The ID of the evaluation dataset that the test case result is for
        evaluation_dataset_version_num: The version number of the evaluation dataset that the
            test case result is for
        test_case_id: The ID of the test case that the test case result is for
        test_case_evaluation_data: A payload representing the data generated by the application
            described by the application spec when evaluated against the test case.
        test_case_evaluation_data_schema: The schema type of the `test_case_evaluation_data`
        result: The payload filled in when the evaluation of this test case result is completed.
            Examine this value to determine how the application performed on the test case.
        completed_at: The time the test case result was completed
        created_at: The time the test case result was created
    """

    id: str
    label_status: EvaluationStatus
    application_spec_id: str
    evaluation_id: str
    evaluation_dataset_id: str
    evaluation_dataset_version_num: str
    test_case_id: str
    test_case_evaluation_data: GenerationTestCaseResultData
    test_case_evaluation_data_schema: TestCaseSchemaType
    created_at: datetime
    result: Optional[Dict[str, Any]] = None
    completed_at: Optional[datetime] = None

    @property
    def status(self):
        warn(
            message="TestCaseResult.status is deprecated. "
                    "Please use TestCaseResult.label_status instead.",
            category=DeprecationWarning,
        )
        return self.label_status


class TestCaseResultSchemaValidator:
    TEST_CASE_RESULT_SCHEMA_TO_DATA_TYPE = {
        TestCaseSchemaType.GENERATION: GenerationTestCaseResultData,
    }

    @classmethod
    def validate(cls, schema_type: TestCaseSchemaType, data: Dict[str, Any]):
        try:
            model = cls.TEST_CASE_RESULT_SCHEMA_TO_DATA_TYPE[schema_type]
            model.from_dict(data)
        except ValidationError as e:
            raise ValueError(
                f"Test case result data does not match schema type {schema_type}"
            ) from e

    @classmethod
    def get_model(cls, schema_type: TestCaseSchemaType):
        return cls.TEST_CASE_RESULT_SCHEMA_TO_DATA_TYPE[schema_type]

    @classmethod
    def dict_to_model(cls, schema_type: TestCaseSchemaType, data: Dict[str, Any]):
        return cls.TEST_CASE_RESULT_SCHEMA_TO_DATA_TYPE[schema_type].from_dict(data)


class TestCaseResultRequest(BaseModel):
    application_spec_id: str
    evaluation_dataset_version_num: int
    test_case_id: str
    test_case_evaluation_data_schema: TestCaseSchemaType
    test_case_evaluation_data: GenerationTestCaseResultData
    result: Optional[Dict[str, Any]] = Field(default_factory=dict)
    account_id: Optional[str] = Field(
        description="Account to create knowledge base in. If you have access to more than one "
                    "account, you must specify an account_id"
    )

    @validator("test_case_evaluation_data")
    @classmethod
    def test_case_evaluation_data_matches_schema_type(cls, test_case_evaluation_data, values):
        schema_type = values.get("test_case_evaluation_data_schema")
        TestCaseResultSchemaValidator.validate(schema_type, test_case_evaluation_data)
        return test_case_evaluation_data
