import datetime
from enum import Enum
from typing import Annotated
from typing import Any
from typing import Literal

import pytest
from amsdal_models.classes.model import Model
from pydantic import BaseModel
from pydantic import BeforeValidator

from amsdal_ml.ml_models.models import MLModel
from amsdal_ml.ml_models.utils import ResponseFormat
from amsdal_ml.ml_retrievers.query_creator import NLQueryCreatorExecutor


class MockLLM(MLModel):
    def __init__(self, response: str, response_format: str | None = None):
        self._response = response
        self._response_format = response_format
        self.async_mode = False

    @property
    def supported_formats(self) -> set[ResponseFormat]:
        return {ResponseFormat.JSON_OBJECT, ResponseFormat.PLAIN_TEXT}

    @property
    def input_role(self) -> str:
        return 'user'

    @property
    def output_role(self) -> str:
        return 'assistant'

    @property
    def tool_role(self) -> str:
        return 'tool'

    @property
    def system_role(self) -> str:
        return 'system'

    @property
    def content_field(self) -> str:
        return 'content'

    @property
    def role_field(self) -> str:
        return 'role'

    @property
    def tool_call_id_field(self) -> str:
        return 'tool_call_id'

    @property
    def tool_name_field(self) -> str:
        return 'name'

    def setup(self) -> None:
        pass

    def teardown(self) -> None:
        pass

    def invoke(self, input_data: Any, **kwargs: Any) -> str:  # type: ignore[override]
        return self._response

    async def ainvoke(self, input_data: Any, **kwargs: Any) -> str:  # type: ignore[override]
        return self._response

    def stream(self, input_data: Any, **kwargs: Any) -> Any:  # type: ignore[override]
        raise NotImplementedError

    async def astream(self, input_data: Any, **kwargs: Any) -> Any:  # type: ignore[override]
        raise NotImplementedError


def _validate_state(v):
    return v


def _validate_category(v):
    return v


class File(BaseModel):
    filename: str
    content_type: str


class StateOptions(str, Enum):
    AL = 'Alabama'
    AK = 'Alaska'
    AZ = 'Arizona'
    CA = 'California'


StateOptionsField = Annotated[StateOptions, BeforeValidator(_validate_state)]


class GuaranteePeriod(str, Enum):
    FIVE = '5'
    TEN = '10'
    FIFTEEN = '15'


class ProductCategory(str, Enum):
    ANNUITY = 'Annuity'
    LIFE = 'Life'
    HEALTH = 'Health'


class Carrier(Model):
    __table_name__ = 'carriers'

    id: int
    name: str
    code: str
    state: StateOptionsField
    category: Annotated[ProductCategory, BeforeValidator(_validate_category)]
    guarantee_period: GuaranteePeriod
    is_active: bool
    premium_amount: float
    coverage_limit: int | None
    effective_date: datetime.date
    expiration_date: datetime.datetime
    description: str | None
    tags: list[str]
    metadata: dict[str, Any]
    file_attachment: File | None
    status: Literal['active', 'inactive', 'pending']
    priority: Literal[1, 2, 3]


class Policy(Model):
    __table_name__ = 'policies'

    id: int
    policy_number: str
    carrier: Carrier | None
    premium: float
    coverage: int


@pytest.mark.asyncio
async def test_nested_fields_detection():
    """Test detection of nested entity fields with various type annotations."""
    executor = NLQueryCreatorExecutor(
        llm=MockLLM(response=''),
        queryset=Policy.objects.all(),
    )

    nested = executor._get_nested_fields()

    assert 'carrier' in nested
    assert nested['carrier'] == 'Carrier'


@pytest.mark.asyncio
async def test_field_type_name_extraction():
    """Test extraction of field type names for various annotations."""
    executor = NLQueryCreatorExecutor(
        llm=MockLLM(response=''),
        queryset=Policy.objects.all(),
    )

    assert executor._get_field_type_name(str) == 'str'

    annotated_field = Carrier.model_fields['category']
    assert executor._get_field_type_name(annotated_field.annotation) == 'ProductCategory'

    carrier_field = Policy.model_fields['carrier']
    assert executor._get_field_type_name(carrier_field.annotation) == 'Carrier'

    status_field = Carrier.model_fields['status']
    assert executor._get_field_type_name(status_field.annotation) == 'Literal'


@pytest.mark.asyncio
async def test_is_nested_entity_detection():
    """Test detection of nested entities in various type annotations."""
    executor = NLQueryCreatorExecutor(
        llm=MockLLM(response=''),
        queryset=Policy.objects.all(),
    )

    assert executor._is_nested_entity(Carrier) is True

    category_field = Carrier.model_fields['category']
    assert executor._is_nested_entity(category_field.annotation) is False

    carrier_field = Policy.model_fields['carrier']
    assert executor._is_nested_entity(carrier_field.annotation) is True

    coverage_field = Carrier.model_fields['coverage_limit']
    assert executor._is_nested_entity(coverage_field.annotation) is False

    assert executor._is_nested_entity(str) is False


@pytest.mark.asyncio
async def test_creation_schema_structure():
    """Test that creation schema is properly generated as JSON schema."""
    executor = NLQueryCreatorExecutor(
        llm=MockLLM(response=''),
        queryset=Carrier.objects.all(),
    )

    schema = executor._build_creation_schema()

    assert isinstance(schema, dict)
    assert '$schema' in schema or 'type' in schema

    assert 'properties' in schema
    properties = schema['properties']

    assert 'id' in properties
    assert 'name' in properties
    assert 'state' in properties
    assert 'category' in properties
    assert 'tags' in properties  # list
    assert 'metadata' in properties  # dict
    assert 'file_attachment' in properties  # File | None

    if 'required' in schema:
        required = schema['required']
        assert 'id' in required
        assert 'name' in required


@pytest.mark.asyncio
async def test_complex_field_types_in_schema():
    """Test that complex field types are represented in JSON schema."""
    executor = NLQueryCreatorExecutor(
        llm=MockLLM(response=''),
        queryset=Carrier.objects.all(),
    )

    schema = executor._build_creation_schema()
    properties = schema['properties']

    state_prop = properties['state']
    assert '$ref' in state_prop
    assert 'enum' in schema['$defs']['StateOptions']

    status_prop = properties['status']
    assert 'enum' in status_prop or 'const' in status_prop

    tags_prop = properties['tags']
    assert tags_prop.get('type') == 'array' or 'items' in tags_prop

    metadata_prop = properties['metadata']
    assert metadata_prop.get('type') == 'object' or 'properties' in metadata_prop


@pytest.mark.asyncio
async def test_union_handling():
    """Test that Union types are handled correctly."""
    executor = NLQueryCreatorExecutor(
        llm=MockLLM(response=''),
        queryset=Policy.objects.all(),
    )

    carrier_field = Policy.model_fields['carrier']
    assert executor._is_nested_entity(carrier_field.annotation) is True

    assert executor._get_field_type_name(carrier_field.annotation) == 'Carrier'
