from __future__ import annotations

import asyncio
import json
import logging
import types
from pathlib import Path
from typing import TYPE_CHECKING
from typing import Any
from typing import ClassVar
from typing import Union
from typing import get_args
from typing import get_origin

from amsdal_models.classes.model import Model
from amsdal_utils.models.data_models.reference import Reference
from amsdal_utils.query.enums import Lookup
from pydantic import BaseModel

from amsdal_ml.ml_models.utils import ResponseFormat
from amsdal_ml.ml_retrievers.query_retriever import NLQueryExecutor
from amsdal_ml.prompts import get_prompt
from amsdal_ml.utils.nl_payload_utils import clean_json_response
from amsdal_ml.utils.reference_utils import build_reference_payload
from amsdal_ml.utils.response_format import select_response_format

if TYPE_CHECKING:
    from amsdal_ml.ml_models.models import MLModel
    from amsdal_ml.ml_retrievers.adapters import RetrieverAdapter
    from amsdal_ml.ml_retrievers.types import AmsdalQuerySet


class ResolutionCandidate(BaseModel):
    field: str
    entity_type: str
    search_keywords: list[str]


class ResolutionPlan(BaseModel):
    candidates: list[ResolutionCandidate]


class ResolvedEntity(BaseModel):
    is_exists: bool
    entity_id: str | None = None
    entity_data: dict[str, Any] | None = None


class ContextResolution(BaseModel):
    resolutions: dict[str, ResolvedEntity]


class ResolutionPlanError(RuntimeError):
    pass


class NestedEntityResolverMixin:
    llm: MLModel
    logger: logging.Logger
    model_class: type[Model]
    resolution_prompt_path: str | Path | None
    llm_response_format: ResponseFormat | None

    @property
    def resolution_response_schema(self) -> dict[str, Any]:
        raise NotImplementedError

    CORE_MODEL_TYPES: ClassVar[set[str]] = {
        'Option',
        'ClassObject',
        'ClassProperty',
        'Fixture',
        'StorageMetadata',
        'Validator',
    }

    def _get_nested_fields(self) -> dict[str, str]:
        nested: dict[str, str] = {}
        model_fields = self.model_class.model_fields

        for field_name, field_obj in model_fields.items():
            field_type = self._get_field_type_name(field_obj.annotation)

            if self._is_nested_entity(field_obj.annotation):
                nested[field_name] = field_type

        return nested

    def _is_nested_entity(self, annotation: Any) -> bool:
        origin = get_origin(annotation)

        if origin is not None and hasattr(origin, '__name__'):
            if origin.__name__ == 'Annotated':
                args = get_args(annotation)
                if args:
                    return self._is_nested_entity(args[0])
            if origin in (Union, types.UnionType):
                args = get_args(annotation)
                for arg in args:
                    if arg is not type(None) and self._is_nested_entity(arg):
                        return True

        try:
            if isinstance(annotation, type) and issubclass(annotation, Model):
                type_name = annotation.__name__
                return type_name not in self.CORE_MODEL_TYPES
        except TypeError:
            pass

        return False

    def _get_field_type_name(self, annotation: Any) -> str:
        origin = get_origin(annotation)

        if origin is not None and hasattr(origin, '__name__'):
            if origin.__name__ == 'Annotated':
                args = get_args(annotation)
                if args:
                    return self._get_field_type_name(args[0])
            if get_origin(annotation) is Reference:
                args = get_args(annotation)
                if args:
                    return self._get_field_type_name(args[0])
            if origin in (Union, types.UnionType):
                args = get_args(annotation)
                for arg in args:
                    if arg is not type(None) and get_origin(arg) is Reference:
                        ref_args = get_args(arg)
                        if ref_args:
                            return self._get_field_type_name(ref_args[0])
                for arg in args:
                    if arg is not type(None) and isinstance(arg, type) and issubclass(arg, Model):
                        return arg.__name__
                for arg in args:
                    if arg is not type(None):
                        return self._get_field_type_name(arg)

        if hasattr(annotation, '__name__'):
            return annotation.__name__

        return str(annotation)

    def _get_entity_class(self, entity_type: str) -> type[Model] | None:
        try:
            from amsdal_ml.utils.model_scanner import get_all_models

            loaded_models = get_all_models()

            for model in loaded_models:
                if model.__name__ == entity_type:
                    return model

            return None
        except Exception:
            return None

    async def _phase_resolution(self, query: str) -> ResolutionPlan:
        nested_fields = self._get_nested_fields()

        if not nested_fields:
            return ResolutionPlan(candidates=[])

        prompt = get_prompt('nl_query_create_resolution', custom_path=self.resolution_prompt_path)
        full_prompt = prompt.render_text(
            query=query,
            nested_fields=json.dumps(nested_fields, indent=2),
        )

        response = await self.llm.ainvoke(
            full_prompt,
            response_format=self.llm_response_format,
            schema=(
                self.resolution_response_schema if self.llm_response_format == ResponseFormat.JSON_SCHEMA else None
            ),
        )

        raw_json = clean_json_response(response)

        try:
            data = json.loads(raw_json)
            return ResolutionPlan(**data)
        except Exception as exc:
            msg = f'Failed to parse resolution plan: {exc}'
            raise ResolutionPlanError(msg) from exc

    async def _phase_retrieval(self, plan: ResolutionPlan) -> ContextResolution:
        logger: logging.Logger = self.logger
        tasks = [self._resolve_candidate(candidate) for candidate in plan.candidates]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        resolutions: dict[str, ResolvedEntity] = {}
        for candidate, result in zip(plan.candidates, results, strict=False):
            if isinstance(result, BaseException):
                logger.error('Error resolving %s: %s', candidate.field, result)
                resolutions[candidate.field] = ResolvedEntity(is_exists=False)
            else:
                resolutions[candidate.field] = result

        logger.debug('Completed context resolution with %s entities', len(resolutions))
        return ContextResolution(resolutions=resolutions)

    async def _resolve_candidate(self, candidate: ResolutionCandidate) -> ResolvedEntity:
        if candidate.entity_type == 'File':
            return ResolvedEntity(is_exists=False)

        entity_cls = self._get_entity_class(candidate.entity_type)

        if entity_cls is None:
            return ResolvedEntity(is_exists=False)

        results = await self._search_entity(entity_cls, candidate.search_keywords)

        if results:
            first_result = results[0]
            entity_object_id = str(first_result.object_id)
            entity_data = first_result.model_dump()

            return ResolvedEntity(
                is_exists=True,
                entity_id=entity_object_id,
                entity_data=entity_data,
            )

        return ResolvedEntity(is_exists=False)

    async def _search_entity(self, entity_cls: type[Model], keywords: list[str]) -> list[Model]:
        logger: logging.Logger = self.logger

        if not keywords:
            return []

        try:
            queryset = entity_cls.objects.all().latest()
            executor = NLQueryExecutor(llm=self.llm, queryset=queryset)
            query = f'find {entity_cls.__name__} matching {" ".join(keywords)}'
            analysis = await executor.analyze(query, preview_limit=0, limit=20)
            documents: list[Model] = await executor.execute(analysis)
            return documents[:20]
        except Exception as exc:
            logger.error('Error searching for %s: %s', entity_cls.__name__, exc)
            return []

    def _select_response_format(self, requested: ResponseFormat | None) -> ResponseFormat:
        return select_response_format(self.llm.supported_formats, requested=requested)

    def _apply_context_resolutions(
        self,
        fields: dict[str, object],
        context: ContextResolution,
    ) -> dict[str, object]:
        nested_fields = self._get_nested_fields()
        if not nested_fields:
            return fields

        resolved = dict(fields)
        resolutions = context.resolutions

        for field_name, entity_type in nested_fields.items():
            if field_name not in resolved:
                continue

            value = resolved[field_name]
            if value is None:
                continue

            if isinstance(value, (Reference, Model)):
                continue

            resolution = resolutions.get(field_name)

            if isinstance(value, dict) and '$ref' in value:
                if resolution is not None and resolution.is_exists and resolution.entity_id is not None:
                    resolved[field_name] = build_reference_payload(entity_type, resolution.entity_id)
                    continue

                raw_ref = value.get('$ref')
                if isinstance(raw_ref, dict):
                    ref_class_name = raw_ref.get('class_name')
                    ref_object_id = raw_ref.get('object_id')
                    if isinstance(ref_class_name, str) and ref_object_id is not None:
                        resolved[field_name] = build_reference_payload(ref_class_name, str(ref_object_id))
                        continue

                continue

            if resolution is not None and resolution.is_exists and resolution.entity_id is not None:
                resolved[field_name] = build_reference_payload(entity_type, resolution.entity_id)

        return resolved


class FilterGenerationMixin:
    llm: MLModel
    base_queryset: AmsdalQuerySet[Model]
    fields: list[str] | None
    fields_info: dict[str, str]
    prompt_path: str | Path | None
    llm_response_format: ResponseFormat | None
    adapter: RetrieverAdapter
    logger: logging.Logger

    async def _build_filters(self, find_query: str) -> dict[str, object]:
        selector = NLQueryExecutor(
            llm=self.llm,
            queryset=self.base_queryset,
            fields=self.fields,
            fields_info=self.fields_info,
            prompt_path=self.prompt_path,
            llm_response_format=self.llm_response_format,
            adapter=self.adapter,
        )

        prompt = get_prompt('nl_query_filter', custom_path=self.prompt_path)
        full_prompt = prompt.render_text(
            schema=json.dumps(selector.search_schema, indent=2),
            query=find_query,
        )

        response = await self.llm.ainvoke(
            full_prompt,
            response_format=self.llm_response_format,
            schema=(selector.response_schema if self.llm_response_format == ResponseFormat.JSON_SCHEMA else None),
        )
        raw_json = clean_json_response(response)
        if hasattr(self, 'logger'):
            self.logger.debug('NLQ filter raw: %s', raw_json)

        conditions = selector.adapter.parse_response(
            raw_json,
            is_schema_based=self.llm_response_format == ResponseFormat.JSON_SCHEMA,
        )

        allowed_lookups = {lookup.value for lookup in Lookup}
        out: dict[str, object] = {}
        for cond in conditions:
            if cond.lookup not in allowed_lookups:
                continue
            out[f'{cond.field}__{cond.lookup}'] = cond.value

        return out
