import json
import logging
from collections.abc import Awaitable
from collections.abc import Callable
from enum import Enum
from pathlib import Path
from typing import Any
from typing import Generic
from typing import Optional
from typing import TypeVar

from amsdal_data.transactions.decorators import async_transaction
from amsdal_models.classes.model import Model
from pydantic import BaseModel
from pydantic import Field
from pydantic import ValidationError

from amsdal_ml.ml_models.models import MLModel
from amsdal_ml.ml_models.utils import ResponseFormat
from amsdal_ml.ml_retrievers.adapters import RetrieverAdapter
from amsdal_ml.ml_retrievers.adapters import get_retriever_adapter
from amsdal_ml.ml_retrievers.exceptions import MultipleRecordsError
from amsdal_ml.ml_retrievers.mixins import ContextResolution
from amsdal_ml.ml_retrievers.mixins import NestedEntityResolverMixin
from amsdal_ml.ml_retrievers.mixins import ResolutionPlan
from amsdal_ml.ml_retrievers.retriever import Document
from amsdal_ml.ml_retrievers.types import AmsdalQuerySet
from amsdal_ml.prompts import get_prompt
from amsdal_ml.utils.dict_utils import deep_merge_dicts
from amsdal_ml.utils.nl_payload_utils import clean_json_response
from amsdal_ml.utils.nl_payload_utils import process_payload_fields
from amsdal_ml.utils.query_utils import serialize_and_clean_record

T = TypeVar('T', bound=Model)


class CreatePayload(BaseModel):
    """Final payload for record creation."""

    fields: dict[str, Any] = Field(default_factory=dict)
    missing_required: list[str] = Field(default_factory=list)


class CreateIntent(str, Enum):
    CREATE = 'CREATE'
    AMBIGUOUS = 'AMBIGUOUS'


class CreateAnalysis(BaseModel):
    query: str
    payloads: list[CreatePayload]
    total: int
    records: list[dict[str, object]]
    intent: CreateIntent | None = None
    target_count: int | None = None
    summary: str | None = None


class NLQueryCreatorExecutor(NestedEntityResolverMixin, Generic[T]):
    """
    Natural language creation interface for Amsdal QuerySets.

    Converts natural language creation requests into database records using a 3-phase pipeline:
    1. Resolution Strategy: Analyze query for nested entities (when required)
    2. Context Retrieval: Search for existing nested entities
    3. Payload Synthesis + Persistence: Generate creation payload and save record(s)

    Supported Field Types:
    - Primitives: str, int, float, bool
    - Enums: Enum subclasses
    - Literals: Literal[...]
    - Dates: datetime, date
    - Files: base64 encoded or URL references

    Unsupported Field Types (will be skipped or require manual handling):
    - Complex nested models (non-core)
    - Dicts: dict[...]
    - Lists/Tuples/Sets of complex types

    Example:
        >>> llm = OpenAIModel()
        >>> llm.setup()
        >>> queryset = Product.objects.all()
        >>> executor = NLQueryCreatorExecutor(llm=llm, queryset=queryset)
        >>> analysis = await executor.analyze("Create product iPhone 16 with price 999")
        >>> results = await executor.execute(analysis)
    """

    _full_prompt: Optional[str]

    def __init__(
        self,
        llm: MLModel,
        queryset: AmsdalQuerySet[T],
        base_defaults: Optional[dict[str, Any]] = None,
        fields: Optional[list[str]] = None,
        fields_info: Optional[dict[str, str]] = None,
        prompt_path: Optional[str | Path] = None,
        resolution_prompt_path: Optional[str | Path] = None,
        synthesis_prompt_path: Optional[str | Path] = None,
        llm_response_format: Optional[ResponseFormat] = None,
        adapter: Optional[RetrieverAdapter] = None,
        on_before_save: Optional[Callable[[Model], Awaitable[None]]] = None,
    ) -> None:
        """
        Initialize NLQueryCreatorExecutor.

        Args:
            llm: ML model instance for query interpretation
            queryset: AMSDAL QuerySet for target model
            base_defaults: Optional dict of default values to apply to all created records
            fields: Optional list of field names to include in creation schema.
                   If None, all model fields are used.
            fields_info: Optional dict mapping field names to descriptions.
                        Takes priority over model field titles.
            prompt_path: Optional path to custom prompt file.
            llm_response_format: Optional desired response format.
            adapter: Optional adapter instance.
            on_before_save: Optional async callback invoked before saving a new record (for authorization checks).
        """
        self.llm = llm
        self.base_queryset = queryset
        self.model_class = queryset.entity
        self.base_defaults = base_defaults or {}
        self.fields = fields
        self.fields_info = fields_info or {}
        self.prompt_path = prompt_path
        self.resolution_prompt_path = resolution_prompt_path or prompt_path
        self.synthesis_prompt_path = synthesis_prompt_path or prompt_path
        self.adapter = adapter or get_retriever_adapter(llm)
        self.llm_response_format = self._select_response_format(llm_response_format)
        self.on_before_save = on_before_save

        self.logger = logging.getLogger(__name__)

        self.creation_schema = self._build_creation_schema()
        self._full_prompt = None

    @property
    def resolution_response_schema(self) -> dict[str, Any]:
        base_schema = ResolutionPlan.model_json_schema()
        return self.adapter.get_response_schema(base_schema)

    async def analyze(self, query: str, *, preview_limit: int = 5) -> CreateAnalysis:
        self.logger.info(f'Starting NLQueryCreator analyze for query: {query}')

        resolution_plan = await self._phase_resolution(query)
        context = await self._phase_retrieval(resolution_plan)
        payloads = await self._phase_synthesis(query, context)

        preview_records: list[dict[str, object]] = []
        for payload in payloads[: max(preview_limit, 0)]:
            merged_fields = {**self.base_defaults, **payload.fields}
            preview_records.append(
                {
                    **merged_fields,
                    'missing_required': payload.missing_required,
                }
            )

        total = len(payloads)
        summary = f'{CreateIntent.CREATE.value}: {total} payload(s)'

        return CreateAnalysis(
            query=query,
            payloads=payloads,
            total=total,
            records=preview_records,
            intent=CreateIntent.CREATE,
            target_count=total,
            summary=summary,
        )

    async def execute(self, analysis: CreateAnalysis) -> list[T]:
        return await self.execute_from_analysis(analysis)  # type: ignore[call-arg, arg-type]

    @async_transaction
    async def execute_from_analysis(self, analysis: CreateAnalysis) -> list[T]:
        created: list[T] = []
        for idx, payload in enumerate(analysis.payloads):
            if payload.missing_required:
                msg = f'Missing required fields (item #{idx + 1}): {", ".join(payload.missing_required)}'
                self.logger.warning(msg)
                raise ValueError(msg)
            record: T = await self._phase_persistence(payload)
            created.append(record)
        return created

    async def execute_one(self, query_or_analysis: str | CreateAnalysis, *, preview_limit: int = 5) -> T:
        if isinstance(query_or_analysis, str):
            analysis = await self.analyze(query_or_analysis, preview_limit=preview_limit)
        else:
            analysis = query_or_analysis

        if analysis.total != 1:
            msg = f'Expected exactly 1 record, got {analysis.total} (preview: {len(analysis.records)})'
            raise MultipleRecordsError(msg, total=analysis.total, records=analysis.records)

        created: list[T] = await self.execute(analysis)
        if len(created) != 1:
            msg = f'Expected exactly 1 record, got {len(created)} (preview: {len(analysis.records)})'
            raise MultipleRecordsError(msg, total=len(created), records=analysis.records)
        return created[0]

    async def _phase_synthesis(self, query: str, context: ContextResolution) -> list[CreatePayload]:
        """
        Phase 3: Payload Synthesis.

        Generate final creation payload based on query and resolved context.
        """
        schema_json = json.dumps(self.creation_schema, indent=2, ensure_ascii=False)
        context_json = json.dumps(context.model_dump(), indent=2, ensure_ascii=False, default=str)
        prompt = get_prompt(
            'nl_query_create_synthesis',
            custom_path=self.synthesis_prompt_path,
        )
        full_prompt = prompt.render_text(
            query=query,
            schema=schema_json,
            context=context_json,
        )
        self._full_prompt = full_prompt

        synthesis_format = self.llm_response_format
        if synthesis_format == ResponseFormat.JSON_SCHEMA:
            synthesis_format = (
                ResponseFormat.JSON_OBJECT
                if ResponseFormat.JSON_OBJECT in self.llm.supported_formats
                else ResponseFormat.PLAIN_TEXT
            )

        response = await self.llm.ainvoke(
            full_prompt,
            response_format=synthesis_format,
        )

        self.logger.debug(f'LLM response for synthesis: {response}')

        raw_json = clean_json_response(response)

        try:
            data = json.loads(raw_json)
            if isinstance(data, list):
                payloads = [CreatePayload(**item) for item in data]
            elif isinstance(data, dict) and 'items' in data and isinstance(data['items'], list):
                payloads = [CreatePayload(**item) for item in data['items']]
            elif isinstance(data, dict) and isinstance(data.get('fields'), list):
                missing_required = data.get('missing_required')
                missing_required_list = missing_required if isinstance(missing_required, list) else []
                payloads = [
                    CreatePayload(fields=item, missing_required=missing_required_list)
                    for item in data['fields']
                    if isinstance(item, dict)
                ]
            else:
                payloads = [CreatePayload(**data)]
        except json.JSONDecodeError as exc:
            msg = f'Failed to decode create payload JSON from LLM: {exc}'
            self.logger.error(msg)
            raise ValueError(msg) from exc
        except ValidationError as exc:
            msg = f'Invalid create payload schema from LLM: {exc}'
            self.logger.error(msg)
            raise ValueError(msg) from exc

        resolved_payloads = []
        for payload in payloads:
            resolved_fields = self._apply_context_resolutions(payload.fields, context)
            resolved_payloads.append(payload.model_copy(update={'fields': resolved_fields}))

        return resolved_payloads

    async def _phase_persistence(self, payload: CreatePayload) -> T:
        """
        Phase 4: Middleware & Persistence.

        Process files, convert references, and save record to database.
        """
        processed_fields = await process_payload_fields(
            payload.fields,
            model_class=self.model_class,
            logger=self.logger,
        )
        self.logger.debug(f'Processed payload fields: {processed_fields}')

        if self.base_defaults and processed_fields:
            all_fields = deep_merge_dicts(self.base_defaults, processed_fields)
        else:
            all_fields = {**self.base_defaults, **processed_fields}
        self.logger.info(f'Final fields for creation: {all_fields}')

        instance = self.model_class(**all_fields)
        if self.on_before_save:
            await self.on_before_save(instance)
        await instance.asave(force_insert=True)

        return instance  # type: ignore[return-value]

    def _build_creation_schema(self) -> dict[str, Any]:
        """Build schema for record creation."""
        return self.model_class.model_json_schema()


class NLQueryCreator(Generic[T]):
    """
    Natural language creation facade for Amsdal QuerySets.

    Provides a `Document`-oriented interface on top of `NLQueryCreatorExecutor`
    for workflows that expect serialized records.

    Example:
        >>> llm = OpenAIModel()
        >>> llm.setup()
        >>> queryset = Product.objects.all()
        >>> creator = NLQueryCreator(llm=llm, queryset=queryset)
        >>> analysis = await creator.analyze("Create product iPhone 16")
        >>> docs = await creator.invoke(analysis)
    """

    def __init__(
        self,
        llm: MLModel,
        queryset: AmsdalQuerySet[T],
        base_defaults: Optional[dict[str, Any]] = None,
        fields: Optional[list[str]] = None,
        fields_info: Optional[dict[str, str]] = None,
        prompt_path: Optional[str | Path] = None,
        llm_response_format: Optional[ResponseFormat] = None,
        adapter: Optional[RetrieverAdapter] = None,
    ) -> None:
        """
        Initialize NLQueryCreator.

        Args:
            llm: ML model instance for query interpretation
            queryset: AMSDAL QuerySet for target model
            base_defaults: Optional dict of default values to apply to all created records
            fields: Optional list of field names to include in creation schema
            fields_info: Optional dict mapping field names to descriptions
            prompt_path: Optional path to custom prompt file
            llm_response_format: Optional desired response format
            adapter: Optional adapter instance
        """
        self.executor = NLQueryCreatorExecutor(
            llm=llm,
            queryset=queryset,
            base_defaults=base_defaults,
            fields=fields,
            fields_info=fields_info,
            prompt_path=prompt_path,
            llm_response_format=llm_response_format,
            adapter=adapter,
        )
        self.logger = logging.getLogger(__name__)

    async def analyze(self, query: str, *, preview_limit: int = 5) -> CreateAnalysis:
        return await self.executor.analyze(query, preview_limit=preview_limit)

    async def invoke(self, analysis: CreateAnalysis) -> list[Document]:
        results: list[T] = await self.executor.execute(analysis)
        documents: list[Document] = []
        for item in results:
            cleaned_data = await serialize_and_clean_record(item)
            documents.append(
                Document(
                    page_content=json.dumps(cleaned_data, ensure_ascii=False),
                    metadata=item.model_dump(),
                )
            )
        return documents

    async def invoke_one(self, query_or_analysis: str | CreateAnalysis, *, preview_limit: int = 5) -> Document:
        item: T = await self.executor.execute_one(query_or_analysis, preview_limit=preview_limit)
        cleaned_data = await serialize_and_clean_record(item)
        return Document(
            page_content=json.dumps(cleaned_data, ensure_ascii=False),
            metadata=item.model_dump(),
        )
