from __future__ import annotations

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 TypeVar

from amsdal_data.transactions.decorators import async_transaction
from amsdal_models.classes.model import Model
from pydantic import BaseModel
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 FilterGenerationMixin
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.nl_payload_utils import clean_json_response
from amsdal_ml.utils.query_utils import apply_related_queryset
from amsdal_ml.utils.query_utils import serialize_and_clean_record
from amsdal_ml.utils.response_format import select_response_format

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


class SplitDeleteQuery(BaseModel):
    find_query: str


class DeleteAnalysis(BaseModel):
    query: str
    find_query: str
    filters: dict[str, object]
    total: int
    records: list[dict[str, object]]
    intent: DeleteIntent | None = None
    target_count: int | None = None
    summary: str | None = None


class DeleteIntent(str, Enum):
    DELETE = 'DELETE'
    AMBIGUOUS = 'AMBIGUOUS'


class InvalidDeleteError(RuntimeError):
    pass


class NLQueryDeleterExecutor(FilterGenerationMixin, Generic[T]):
    """
    Natural language deletion interface for Amsdal QuerySets.

    Flow:
    1. Split: LLM derives a filter query from the delete request.
    2. Analysis: Preview matching records.
    3. Execution: Delete the matched records and return them.

    Example:
        >>> llm = OpenAIModel()
        >>> llm.setup()
        >>> queryset = Product.objects.all()
        >>> executor = NLQueryDeleterExecutor(llm=llm, queryset=queryset)
        >>> analysis = await executor.analyze("delete discontinued products")
        >>> deleted = await executor.execute(analysis)
    """

    _full_prompt: str | None

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

    def __init__(
        self,
        llm: MLModel,
        queryset: AmsdalQuerySet[T],
        *,
        fields: list[str] | None = None,
        fields_info: dict[str, str] | None = None,
        prompt_path: str | Path | None = None,
        llm_response_format: ResponseFormat | None = None,
        adapter: RetrieverAdapter | None = None,
        on_before_delete: Callable[[Model], Awaitable[None]] | None = None,
    ) -> None:
        self.llm = llm
        self.base_queryset = queryset  # type: ignore[assignment]
        self.base_filters: dict[str, object] = {'_metadata__is_deleted': False}
        self.model_class = queryset.entity
        self.fields = fields
        self.fields_info = fields_info or {}
        self.prompt_path = prompt_path
        self.adapter = adapter or get_retriever_adapter(llm)
        self.llm_response_format = self._select_response_format(llm_response_format)
        self.on_before_delete = on_before_delete

        self.logger = logging.getLogger(__name__)
        self._full_prompt = None

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

    async def execute_one(self, query_or_analysis: str | DeleteAnalysis, *, 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)

        queryset = self.base_queryset.latest().filter(**self.base_filters)
        queryset = apply_related_queryset(queryset, analysis.filters)
        queryset = queryset.filter(**analysis.filters)
        targets: list[T] = await queryset[0:2].aexecute()  # type: ignore[attr-defined]
        if len(targets) != 1:
            msg = f'Expected exactly 1 record, got {len(targets)} (preview: {len(analysis.records)})'
            raise MultipleRecordsError(msg, total=len(targets), records=analysis.records)

        target = targets[0]
        if self.on_before_delete:
            await self.on_before_delete(target)
        await target.adelete()
        return target

    async def analyze(self, query: str, *, preview_limit: int = 5) -> DeleteAnalysis:
        if preview_limit < 0:
            msg = 'preview_limit must be >= 0'
            raise ValueError(msg)

        find_query = await self._split_query(query)
        filters = await self._build_filters(find_query)
        queryset = self.base_queryset.latest().filter(**self.base_filters)
        queryset = apply_related_queryset(queryset, filters)
        queryset = queryset.filter(**filters)
        total = await queryset.count().aexecute()  # type: ignore[attr-defined]

        if preview_limit == 0:
            summary = f'{DeleteIntent.DELETE.value}: {total} target(s)'
            return DeleteAnalysis(
                query=query,
                find_query=find_query,
                filters=filters,
                total=total,
                records=[],
                intent=DeleteIntent.DELETE,
                target_count=total,
                summary=summary,
            )

        preview_qs = queryset[0:preview_limit]
        preview_models: list[T] = await preview_qs.aexecute()  # type: ignore[attr-defined]
        records = [await serialize_and_clean_record(item) for item in preview_models]
        summary = f'{DeleteIntent.DELETE.value}: {total} target(s)'
        return DeleteAnalysis(
            query=query,
            find_query=find_query,
            filters=filters,
            total=total,
            records=records,
            intent=DeleteIntent.DELETE,
            target_count=total,
            summary=summary,
        )

    @async_transaction
    async def execute_from_analysis(self, analysis: DeleteAnalysis) -> list[T]:
        queryset = self.base_queryset.latest().filter(**self.base_filters)
        queryset = apply_related_queryset(queryset, analysis.filters)
        queryset = queryset.filter(**analysis.filters)
        targets: list[T] = await queryset.aexecute()  # type: ignore[attr-defined]
        deleted: list[T] = []

        for target in targets:
            if self.on_before_delete:
                await self.on_before_delete(target)
            await target.adelete()
            deleted.append(target)

        return deleted

    async def _split_query(self, query: str) -> str:
        prompt = get_prompt('nl_query_delete_split', custom_path=self.prompt_path)
        full_prompt = prompt.render_text(query=query)
        self._full_prompt = full_prompt

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

        raw_json = clean_json_response(response)
        try:
            data = json.loads(raw_json)
            split = SplitDeleteQuery(**data)
        except json.JSONDecodeError as exc:
            msg = f'Failed to decode delete split JSON from LLM: {exc}'
            raise InvalidDeleteError(msg) from exc
        except ValidationError as exc:
            msg = f'Invalid delete split JSON schema from LLM: {exc}'
            raise InvalidDeleteError(msg) from exc

        find_query = split.find_query.strip()
        if not find_query:
            msg = 'Split query resulted in empty find_query'
            raise InvalidDeleteError(msg)
        return find_query

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


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

    Wraps `NLQueryDeleterExecutor` and returns deleted records as `Document` objects.

    Example:
        >>> llm = OpenAIModel()
        >>> llm.setup()
        >>> queryset = Product.objects.all()
        >>> deleter = NLQueryDeleter(llm=llm, queryset=queryset)
        >>> analysis = await deleter.analyze("delete discontinued products")
        >>> docs = await deleter.invoke(analysis)
    """

    def __init__(
        self,
        llm: MLModel,
        queryset: AmsdalQuerySet[T],
        *,
        fields: list[str] | None = None,
        fields_info: dict[str, str] | None = None,
        prompt_path: str | Path | None = None,
        llm_response_format: ResponseFormat | None = None,
        adapter: RetrieverAdapter | None = None,
    ) -> None:
        self.executor = NLQueryDeleterExecutor(
            llm=llm,
            queryset=queryset,
            fields=fields,
            fields_info=fields_info,
            prompt_path=prompt_path,
            llm_response_format=llm_response_format,
            adapter=adapter,
        )

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

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

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