from __future__ import annotations

import asyncio
import inspect
import json
import logging
import types
from collections.abc import AsyncIterator
from collections.abc import Awaitable
from collections.abc import Callable
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Annotated
from typing import Any
from typing import Generic
from typing import TypeVar
from typing import Union
from typing import get_args
from typing import get_origin

from amsdal_data.transactions.decorators import async_transaction
from amsdal_models.classes.model import Model
from amsdal_models.classes.model import TypeModel
from amsdal_utils.models.data_models.reference import Reference
from amsdal_utils.models.utils.reference_builders import build_reference
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 FilterGenerationMixin
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.model_scanner import get_all_models
from amsdal_ml.utils.model_serialization import SerializeMode
from amsdal_ml.utils.model_serialization import normalize_value_for_diff
from amsdal_ml.utils.model_serialization import serialize_model
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 apply_related_queryset
from amsdal_ml.utils.query_utils import serialize_and_clean_record
from amsdal_ml.utils.reference_utils import build_reference_payload

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


@dataclass
class _RelatedUpdateRequest:
    executor: NLQueryUpdaterExecutor[Model]
    entity: Model
    payload: UpdatePayload
    context: ContextResolution
    update_query: str

    async def apply(self) -> None:
        await self.executor.apply_payload(self.entity, self.payload, self.context, self.update_query)


class UpdatePayload(BaseModel):
    fields: dict[str, object] = Field(default_factory=dict)
    missing_required: list[str] = Field(default_factory=list)


class UpdateIntent(str, Enum):
    UPDATE_RELATED = 'UPDATE_RELATED'
    REASSIGN_FK = 'REASSIGN_FK'
    UPDATE_DIRECT = 'UPDATE_DIRECT'
    MIXED = 'MIXED'
    AMBIGUOUS = 'AMBIGUOUS'


class UpdateFieldType(str, Enum):
    DIRECT = 'direct'
    FK = 'fk'
    JSONB = 'jsonb'


class UpdateAnalysis(BaseModel):
    query: str
    find_query: str
    update_query: str
    filters: dict[str, object]
    total: int
    records: list[dict[str, object]]
    intent: UpdateIntent | None = None
    target_count: int | None = None
    target_updates: list[EntityUpdatePreview] = Field(default_factory=list)
    related_updates: list[RelatedUpdatePreview] = Field(default_factory=list)
    referenced_by: dict[str, ReferencedByInfo] = Field(default_factory=dict)
    alternative_interpretations: list[UpdateAnalysis] = Field(default_factory=list)
    summary: str | None = None


class UpdateChange(BaseModel):
    field_path: str
    current_value: object
    updated_value: object
    field_type: UpdateFieldType


class EntityUpdatePreview(BaseModel):
    entity_type: str
    entity_id: str | int
    entity_repr: str
    changes: list[UpdateChange]


class ReferencedBySample(BaseModel):
    entity_id: str | int
    entity_repr: str


class ReferencedByInfo(BaseModel):
    count: int
    reference_field: str
    sample: list[ReferencedBySample]


class RelatedUpdatePreview(BaseModel):
    entity_update: EntityUpdatePreview
    referenced_by: dict[str, ReferencedByInfo]


class SplitQuery(BaseModel):
    find_query: str
    update_query: str


class InvalidPatchError(RuntimeError):
    pass


class AmbiguousUpdateIntentError(RuntimeError):
    pass


class NLQueryUpdaterExecutor(FilterGenerationMixin, NestedEntityResolverMixin, Generic[T]):
    """
    Natural language update interface for Amsdal QuerySets.

    Flow:
    1. Split: LLM derives find + update instructions.
    2. Resolution: Optional lookup for nested references.
    3. Patch: Synthesize an update payload and apply it to matched records.

    Example:
        >>> llm = OpenAIModel()
        >>> llm.setup()
        >>> queryset = Product.objects.all()
        >>> executor = NLQueryUpdaterExecutor(llm=llm, queryset=queryset)
        >>> analysis = await executor.analyze("update products set price to 99 where brand is ACME")
        >>> updated = await executor.execute(analysis)
    """

    _full_prompt: str | None
    RELATED_SAMPLE_LIMIT = 5

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

    @property
    def split_response_schema(self) -> dict[str, Any]:
        base_schema = SplitQuery.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_save: 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.resolution_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_save = on_before_save

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

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

    async def execute_one(self, query_or_analysis: str | UpdateAnalysis, *, 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]
        payload, context = await self._build_payload(target, analysis.update_query)
        return await self._apply_payload(target, payload, context, analysis.update_query)

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

        find_query, update_query = await self._split_query(query)
        self.logger.debug('NLQ update split: query=%s find=%s update=%s', query, find_query, update_query)
        filters = await self._build_filters(find_query)
        self.logger.debug('NLQ update filters: %s', filters)
        try:
            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]
        except Exception as exc:  # noqa: BLE001
            self.logger.exception('NLQ update count failed: %s', exc)
            raise
        self.logger.debug('NLQ update total: %s', total)

        if preview_limit == 0:
            return UpdateAnalysis(
                query=query,
                find_query=find_query,
                update_query=update_query,
                filters=filters,
                total=total,
                records=[],
                intent=None,
                target_count=total,
                target_updates=[],
                related_updates=[],
                referenced_by={},
                summary=None,
            )

        preview_qs = queryset[0:preview_limit]
        preview_models: list[T] = await preview_qs.aexecute()  # type: ignore[attr-defined]
        self.logger.debug('NLQ update preview count: %s', len(preview_models))
        records = [await serialize_and_clean_record(item) for item in preview_models]

        tasks = [self._build_update_preview(target, update_query) for target in preview_models]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        target_updates: list[EntityUpdatePreview] = []
        for target, result in zip(preview_models, results, strict=False):
            if isinstance(result, BaseException):
                self.logger.error('Failed to build update preview for %s: %s', target, result)
                continue
            target_updates.append(result)

        related_updates = await self._build_related_update_previews(preview_models, update_query, target_updates)
        self.logger.debug('NLQ update target updates: %s', len(target_updates))
        self.logger.debug('NLQ update related updates: %s', len(related_updates))
        intent = self._infer_intent(target_updates, related_updates)
        referenced_by: dict[str, ReferencedByInfo] = {}
        if len(preview_models) == 1:
            referenced_by = await self._build_referenced_by(preview_models[0])
        summary = self._build_summary(intent, total, target_updates, related_updates)
        self.logger.debug('NLQ update intent: %s summary=%s', intent, summary)
        if intent == UpdateIntent.AMBIGUOUS:
            msg = 'Ambiguous update intent. Please clarify your request.'
            raise AmbiguousUpdateIntentError(msg)
        return UpdateAnalysis(
            query=query,
            find_query=find_query,
            update_query=update_query,
            filters=filters,
            total=total,
            records=records,
            intent=intent,
            target_count=total,
            target_updates=target_updates,
            related_updates=related_updates,
            referenced_by=referenced_by,
            summary=summary,
        )

    @async_transaction
    async def execute_from_analysis(self, analysis: UpdateAnalysis) -> list[T]:
        queryset = self.base_queryset.latest().filter(**self.base_filters)
        queryset = apply_related_queryset(queryset, analysis.filters)
        queryset = queryset.filter(**analysis.filters)
        self.logger.debug('NLQ update execute filters: %s', analysis.filters)
        try:
            targets: list[T] = await queryset.aexecute()  # type: ignore[attr-defined]
        except Exception as exc:  # noqa: BLE001
            self.logger.exception('NLQ update execute failed: %s', exc)
            raise
        updated: list[T] = []

        related_updates = await self._build_related_update_requests(targets, analysis.update_query)
        for related_request in related_updates.values():
            await related_request.apply()

        for target in targets:
            payload, context = await self._build_payload(target, analysis.update_query)
            if not await self._has_changes(target, payload, context, analysis.update_query):
                continue
            applied = await self._apply_payload(target, payload, context, analysis.update_query)
            updated.append(applied)

        return updated

    async def _split_query(self, query: str) -> tuple[str, str]:
        prompt = get_prompt('nl_query_update_split', custom_path=self.prompt_path)
        full_prompt = prompt.render_text(query=query)
        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 = SplitQuery(**data)
        except json.JSONDecodeError as exc:
            msg = f'Failed to decode split JSON from LLM: {exc}'
            raise InvalidPatchError(msg) from exc
        except ValidationError as exc:
            msg = f'Invalid split JSON schema from LLM: {exc}'
            raise InvalidPatchError(msg) from exc

        find_query = split.find_query.strip()
        update_query = split.update_query.strip()
        if not find_query or not update_query:
            msg = 'Split query resulted in empty find/update query'
            raise InvalidPatchError(msg)
        return find_query, update_query

    async def _build_payload(self, target: T, update_query: str) -> tuple[UpdatePayload, ContextResolution]:
        resolution_plan = await self._phase_resolution(update_query)
        context = await self._phase_retrieval(resolution_plan)
        current = await self._serialize_model(target, mode=SerializeMode.PROMPT)
        current_json = json.dumps(current, ensure_ascii=False, indent=2, default=str)
        schema_json = json.dumps(self.model_class.model_json_schema(), ensure_ascii=False, indent=2)
        context_json = json.dumps(context.model_dump(), ensure_ascii=False, indent=2, default=str)

        prompt = get_prompt('nl_query_update_patch', custom_path=self.prompt_path)
        full_prompt = prompt.render_text(
            query=update_query,
            current=current_json,
            schema=schema_json,
            context=context_json,
        )
        self._full_prompt = full_prompt

        payload_format = self.llm_response_format
        if payload_format == ResponseFormat.JSON_SCHEMA:
            payload_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=payload_format,
        )

        raw_json = clean_json_response(response)

        try:
            data = json.loads(raw_json)
            payload = UpdatePayload(**data)
        except json.JSONDecodeError as exc:
            msg = f'Failed to decode update payload JSON from LLM: {exc}'
            raise InvalidPatchError(msg) from exc
        except ValidationError as exc:
            msg = f'Invalid update payload schema from LLM: {exc}'
            raise InvalidPatchError(msg) from exc

        if payload.missing_required:
            msg = f'Missing required fields: {", ".join(payload.missing_required)}'
            raise InvalidPatchError(msg)

        if not payload.fields:
            msg = 'Update payload resulted in empty fields'
            raise InvalidPatchError(msg)

        return payload, context

    async def apply_payload(
        self,
        target: T,
        payload: UpdatePayload,
        context: ContextResolution,
        update_query: str,
    ) -> T:
        return await self._apply_payload(target, payload, context, update_query)

    async def _build_update_preview(self, target: T, update_query: str) -> EntityUpdatePreview:
        payload, context = await self._build_payload(target, update_query)
        self.logger.debug('NLQ update payload fields: %s', list(payload.fields.keys()))
        processed_fields = await self._prepare_payload_fields(target, payload, context, update_query)
        current_values = await self._serialize_model(target, mode=SerializeMode.DIFF)
        changes = self._collect_changes(current_values, processed_fields)
        self.logger.debug('NLQ update change count: %s', len(changes))
        entity_id = self._get_entity_id(target)
        entity_repr = await self._get_entity_repr(target)
        return EntityUpdatePreview(
            entity_type=self.model_class.__name__,
            entity_id=entity_id,
            entity_repr=entity_repr,
            changes=changes,
        )

    async def _prepare_payload_fields(
        self,
        target: T,
        payload: UpdatePayload,
        context: ContextResolution,
        update_query: str,
    ) -> dict[str, object]:
        current = await self._serialize_model(target, mode=SerializeMode.PAYLOAD)
        model_fields = self.model_class.model_fields

        raw_fields: dict[str, object] = {**current, **payload.fields}
        raw_fields = self._merge_typemodel_payloads(current, payload.fields, raw_fields)
        filtered_fields: dict[str, object] = {name: value for name, value in raw_fields.items() if name in model_fields}
        filtered_fields = self._filter_unmentioned_fk_fields(filtered_fields, update_query)

        resolved_fields = await self._resolve_nested_fields(filtered_fields, context, update_query)
        processed_fields = await process_payload_fields(
            resolved_fields,
            model_class=self.model_class,
            logger=self.logger,
        )

        return processed_fields

    def _merge_typemodel_payloads(
        self,
        current: dict[str, object],
        payload_fields: dict[str, object],
        merged: dict[str, object],
    ) -> dict[str, object]:
        out = dict(merged)
        for field_name, patch_value in payload_fields.items():
            model_field = self.model_class.model_fields.get(field_name)
            if model_field is None:
                continue

            typemodel_cls = self._get_typemodel_class(model_field.annotation)
            if typemodel_cls is None:
                continue

            base_value = current.get(field_name)
            if not isinstance(base_value, dict) or not isinstance(patch_value, dict):
                continue

            out[field_name] = deep_merge_dicts(base_value, patch_value)

        return out

    async def _has_changes(
        self,
        target: T,
        payload: UpdatePayload,
        context: ContextResolution,
        update_query: str,
    ) -> bool:
        processed_fields = await self._prepare_payload_fields(target, payload, context, update_query)
        current_values = await self._serialize_model(target, mode=SerializeMode.DIFF)
        return bool(self._collect_changes(current_values, processed_fields))

    async def _serialize_model(self, target: T, *, mode: SerializeMode) -> dict[str, object]:
        return await serialize_model(target, mode=mode, model_class=self.model_class)

    def _collect_changes(
        self,
        current: dict[str, object],
        updated: dict[str, object],
    ) -> list[UpdateChange]:
        changes: list[UpdateChange] = []
        model_fields = self.model_class.model_fields

        for field_name, current_value in current.items():
            if field_name not in updated:
                continue

            updated_value = normalize_value_for_diff(updated[field_name])
            model_field = model_fields.get(field_name)
            if model_field is None:
                continue

            field_kind = self._get_field_kind(model_field.annotation)
            if (
                field_kind == UpdateFieldType.JSONB
                and isinstance(current_value, dict)
                and isinstance(updated_value, dict)
            ):
                self._diff_jsonb(field_name, current_value, updated_value, changes)
                continue

            if current_value != updated_value:
                changes.append(
                    UpdateChange(
                        field_path=field_name,
                        current_value=current_value,
                        updated_value=updated_value,
                        field_type=field_kind,
                    )
                )

        return changes

    def _diff_jsonb(
        self,
        prefix: str,
        current: dict[str, object],
        updated: dict[str, object],
        changes: list[UpdateChange],
    ) -> None:
        keys = set(current.keys()) | set(updated.keys())
        for key in keys:
            current_value = current.get(key)
            updated_value = updated.get(key)
            path = f'{prefix}.{key}'

            if isinstance(current_value, dict) and isinstance(updated_value, dict):
                self._diff_jsonb(path, current_value, updated_value, changes)
                continue

            if current_value != updated_value:
                changes.append(
                    UpdateChange(
                        field_path=path,
                        current_value=current_value,
                        updated_value=updated_value,
                        field_type=UpdateFieldType.JSONB,
                    )
                )

    def _get_field_kind(self, annotation: Any) -> UpdateFieldType:
        if self._get_model_class(annotation) is not None:
            return UpdateFieldType.FK
        if self._get_typemodel_class(annotation) is not None:
            return UpdateFieldType.JSONB
        return UpdateFieldType.DIRECT

    def _get_model_class(self, annotation: Any) -> type[Model] | None:
        for typ in self._iter_annotation_types(annotation):
            try:
                if isinstance(typ, type) and issubclass(typ, Model):
                    return typ
            except TypeError:
                continue
        return None

    def _get_typemodel_class(self, annotation: Any) -> type[BaseModel] | None:
        for typ in self._iter_annotation_types(annotation):
            try:
                if isinstance(typ, type) and issubclass(typ, TypeModel):
                    return typ
                if isinstance(typ, type) and issubclass(typ, BaseModel) and not issubclass(typ, Model):
                    return typ
            except TypeError:
                continue
        return None

    def _iter_annotation_types(self, annotation: Any) -> list[Any]:
        origin = get_origin(annotation)

        if origin is Annotated:
            args = get_args(annotation)
            if args:
                return self._iter_annotation_types(args[0])

        if origin is Reference:
            args = get_args(annotation)
            if args:
                return self._iter_annotation_types(args[0])

        if origin in (Union, types.UnionType):
            collected_types: list[Any] = []
            for arg in get_args(annotation):
                if arg is type(None):
                    continue
                collected_types.extend(self._iter_annotation_types(arg))
            return collected_types

        return [annotation]

    def _infer_intent(
        self,
        target_updates: list[EntityUpdatePreview],
        related_updates: list[RelatedUpdatePreview],
    ) -> UpdateIntent | None:
        if not target_updates and not related_updates:
            return UpdateIntent.AMBIGUOUS

        has_fk = False
        has_direct = False
        has_jsonb = False
        has_related = bool(related_updates)

        for update in target_updates:
            for change in update.changes:
                if change.field_type == UpdateFieldType.FK:
                    has_fk = True
                elif change.field_type == UpdateFieldType.JSONB:
                    has_jsonb = True
                else:
                    has_direct = True

        self.logger.debug(
            'NLQ update intent flags: has_fk=%s has_direct=%s has_jsonb=%s has_related=%s',
            has_fk,
            has_direct,
            has_jsonb,
            has_related,
        )

        if has_related and (has_direct or has_jsonb or has_fk):
            return UpdateIntent.MIXED
        if has_related:
            return UpdateIntent.UPDATE_RELATED
        if has_fk and (has_direct or has_jsonb):
            return UpdateIntent.MIXED
        if has_fk:
            return UpdateIntent.REASSIGN_FK
        if has_jsonb or has_direct:
            return UpdateIntent.UPDATE_DIRECT
        return UpdateIntent.AMBIGUOUS

    def _build_summary(
        self,
        intent: UpdateIntent | None,
        total: int,
        target_updates: list[EntityUpdatePreview],
        related_updates: list[RelatedUpdatePreview],
    ) -> str | None:
        if intent is None:
            return None

        change_count = sum(len(update.changes) for update in target_updates)
        related_count = len(related_updates)
        if change_count == 0 and related_count == 0:
            return f'{intent.value}: no changes detected'
        return f'{intent.value}: {total} target(s), {change_count} change(s), {related_count} related update(s)'

    async def _get_entity_repr(self, target: T) -> str:
        display_name = target.display_name
        if inspect.isawaitable(display_name):
            return await display_name
        return display_name

    def _get_entity_id(self, target: T) -> str | int:
        return str(target.object_id)

    async def _build_related_update_previews(
        self,
        targets: list[T],
        update_query: str,
        target_updates: list[EntityUpdatePreview],
    ) -> list[RelatedUpdatePreview]:
        updates: list[RelatedUpdatePreview] = []
        preview_by_id = {update.entity_id: update for update in target_updates}

        for target in targets:
            target_preview = preview_by_id.get(self._get_entity_id(target))
            updates.extend(await self._build_related_updates_for_target(target, update_query, target_preview))

        return updates

    async def _build_related_updates_for_target(
        self,
        target: T,
        update_query: str,
        target_preview: EntityUpdatePreview | None,
    ) -> list[RelatedUpdatePreview]:
        related_updates: list[RelatedUpdatePreview] = []
        skip_fields: set[str] = set()
        if target_preview is not None:
            for change in target_preview.changes:
                if change.field_type == UpdateFieldType.FK:
                    skip_fields.add(change.field_path)

        async for (
            _field_name,
            _model_cls,
            related_entity,
            related_executor,
            related_query,
        ) in self._iter_related_update_contexts(target, update_query, skip_fields=skip_fields):
            related_preview = await related_executor._build_update_preview(  # noqa: SLF001
                related_entity,
                related_query,
            )
            if not related_preview.changes:
                continue

            referenced_by = await self._build_referenced_by(related_entity)
            related_updates.append(
                RelatedUpdatePreview(
                    entity_update=related_preview,
                    referenced_by=referenced_by,
                )
            )

        return related_updates

    async def _iter_related_update_contexts(
        self,
        target: T,
        update_query: str,
        *,
        skip_fields: set[str] | None = None,
    ) -> AsyncIterator[tuple[str, type[Model], Model, NLQueryUpdaterExecutor[Model], str]]:
        skip = skip_fields or set()
        for field_name, model_cls in self._iter_fk_fields(self.model_class):
            has_intent = self._has_related_update_intent(update_query, field_name, model_cls)
            self.logger.debug(
                'NLQ update related intent: field=%s model=%s intent=%s',
                field_name,
                model_cls.__name__,
                has_intent,
            )
            if not has_intent or field_name in skip:
                continue

            related_entity = await self._load_related_entity(target, field_name, model_cls)
            if related_entity is None:
                continue

            related_executor = self._build_related_executor(model_cls)
            related_query = await self._normalize_related_update_query(update_query, field_name, model_cls)
            yield field_name, model_cls, related_entity, related_executor, related_query

    def _has_related_update_intent(
        self,
        update_query: str,
        fk_field_name: str,
        model_cls: type[Model],
    ) -> bool:
        query_lower = update_query.lower()
        fk_field = self.model_class.model_fields.get(fk_field_name)
        fk_title = fk_field.title.lower() if fk_field and fk_field.title else ''
        mentions_fk = fk_field_name.lower() in query_lower or (fk_title and fk_title in query_lower)
        if not mentions_fk:
            return False

        for field_name, field in model_cls.model_fields.items():
            if field_name.lower() in query_lower:
                return True
            title = field.title
            if isinstance(title, str) and title.lower() in query_lower:
                return True
        return False

    async def _build_referenced_by(self, entity: Model) -> dict[str, ReferencedByInfo]:
        referenced_by: dict[str, ReferencedByInfo] = {}
        entity_cls = entity.__class__
        for model_cls in get_all_models():
            for field_name, field in model_cls.model_fields.items():
                field_model = self._get_model_class(field.annotation)
                field_reference = self._is_reference_field(field.annotation)
                if field_model is None and not field_reference:
                    continue

                if field_model is not None and field_model is not entity_cls:
                    continue

                if field_reference and not self._is_reference_target(field.annotation, entity_cls):
                    continue

                use_reference = field_reference
                filters = self._build_reference_filter(field_name, entity, use_reference=use_reference)
                if filters is None:
                    continue

                queryset = model_cls.objects.all().latest().filter(**self.base_filters).filter(**filters)
                total = await queryset.count().aexecute()
                if total == 0:
                    continue
                sample_qs = queryset[0 : self.RELATED_SAMPLE_LIMIT]
                sample_models: list[Model] = await sample_qs.aexecute()

                samples = []
                for item in sample_models:
                    entity_repr = await self._get_entity_repr(item)  # type: ignore[arg-type]
                    samples.append(
                        ReferencedBySample(
                            entity_id=str(item.object_id),
                            entity_repr=entity_repr,
                        ),
                    )

                referenced_by[model_cls.__name__] = ReferencedByInfo(
                    count=total,
                    reference_field=field_name,
                    sample=samples,
                )

        return referenced_by

    def _is_reference_field(self, annotation: Any) -> bool:
        origin = get_origin(annotation)
        if origin is Reference:
            return True
        if origin in (Union, types.UnionType):
            return any(self._is_reference_field(arg) for arg in get_args(annotation) if arg is not type(None))
        return False

    def _is_reference_target(self, annotation: Any, target: type[Model]) -> bool:
        origin = get_origin(annotation)
        if origin is Reference:
            args = get_args(annotation)
            return bool(args and args[0] is target)
        if origin in (Union, types.UnionType):
            return any(self._is_reference_target(arg, target) for arg in get_args(annotation) if arg is not type(None))
        return False

    def _build_reference_filter(
        self,
        field_name: str,
        entity: Model,
        *,
        use_reference: bool,
    ) -> dict[str, object] | None:
        if use_reference:
            ref = self._build_reference(entity)
            if ref is None:
                return None
            return {f'{field_name}__eq': ref}

        return {f'{field_name}__eq': entity}

    def _build_reference(self, entity: Model) -> Reference | None:
        object_id = entity.object_id
        class_name = entity.__class__.__name__
        return build_reference(class_name=class_name, object_id=object_id)

    def _iter_fk_fields(self, model_cls: type[Model]) -> list[tuple[str, type[Model]]]:
        fields: list[tuple[str, type[Model]]] = []
        for field_name, field in model_cls.model_fields.items():
            field_model = self._get_model_class(field.annotation)
            if field_model is not None:
                fields.append((field_name, field_model))
        return fields

    async def _load_related_entity(self, target: T, field_name: str, model_cls: type[Model]) -> Model | None:
        raw_value = getattr(target, field_name)
        if inspect.isawaitable(raw_value):
            raw_value = await raw_value

        if isinstance(raw_value, Model):
            return raw_value

        if isinstance(raw_value, Reference):
            ref = raw_value.ref
            object_id = ref.get('object_id') if isinstance(ref, dict) else None
            if object_id is None:
                return None
            queryset = model_cls.objects.all().latest().filter(**self.base_filters).filter(object_id__eq=object_id)
            results: list[Model] = await queryset.aexecute()
            return results[0] if results else None

        return None

    def _build_related_executor(self, model_cls: type[Model]) -> NLQueryUpdaterExecutor[Model]:
        queryset = model_cls.objects.all()
        return NLQueryUpdaterExecutor(
            llm=self.llm,
            queryset=queryset,
            fields=self.fields,
            fields_info=self.fields_info,
            prompt_path=self.prompt_path,
            llm_response_format=self.llm_response_format,
            adapter=self.adapter,
        )

    async def _build_related_update_requests(
        self,
        targets: list[T],
        update_query: str,
    ) -> dict[str, _RelatedUpdateRequest]:
        sem = asyncio.Semaphore(10)

        async def _run_with_sem(coro):
            async with sem:
                return await coro

        context_tasks = [
            _run_with_sem(self._collect_related_contexts_for_target(target, update_query)) for target in targets
        ]
        context_results = await asyncio.gather(*context_tasks, return_exceptions=True)

        payload_tasks = []
        task_metadata = []

        for target_idx, result in enumerate(context_results):
            if isinstance(result, BaseException):
                self.logger.error('Failed to collect related contexts for target %s: %s', targets[target_idx], result)
                continue

            for context in result:
                _, model_cls, related_entity, related_executor, related_query = context
                task = _run_with_sem(
                    self._build_payload_and_check_changes(related_executor, related_entity, related_query)
                )
                payload_tasks.append(task)
                task_metadata.append((model_cls, related_entity, related_executor, related_query))

        if not payload_tasks:
            return {}

        payload_results = await asyncio.gather(*payload_tasks, return_exceptions=True)

        requests: dict[str, _RelatedUpdateRequest] = {}
        for (model_cls, related_entity, related_executor, related_query), result in zip(
            task_metadata,
            payload_results,
            strict=False,
        ):
            if isinstance(result, BaseException):
                self.logger.error('Failed to build payload for %s: %s', related_entity, result)
                continue

            payload, context, has_changes = result
            if not has_changes:
                continue

            key = f'{model_cls.__name__}:{related_entity.object_id}'
            if key in requests:
                continue
            requests[key] = _RelatedUpdateRequest(
                executor=related_executor,
                entity=related_entity,
                payload=payload,
                context=context,
                update_query=related_query,
            )

        return requests

    async def _collect_related_contexts_for_target(
        self,
        target: T,
        update_query: str,
    ) -> list[tuple[str, type[Model], Model, NLQueryUpdaterExecutor[Model], str]]:
        contexts = []
        async for context in self._iter_related_update_contexts(target, update_query):
            contexts.append(context)
        return contexts

    async def _build_payload_and_check_changes(
        self,
        executor: NLQueryUpdaterExecutor[Model],
        entity: Model,
        query: str,
    ) -> tuple[UpdatePayload, ContextResolution, bool]:
        payload, context = await executor._build_payload(entity, query)  # noqa: SLF001
        has_changes = await executor._has_changes(entity, payload, context, query)  # noqa: SLF001
        return payload, context, has_changes

    async def _apply_payload(
        self,
        target: T,
        payload: UpdatePayload,
        context: ContextResolution,
        update_query: str,
    ) -> T:
        model_fields = self.model_class.model_fields
        processed_fields = await self._prepare_payload_fields(target, payload, context, update_query)

        for name, value in processed_fields.items():
            model_field = model_fields.get(name)
            if model_field is None:
                msg = f'Unknown field in fields: {name}'
                raise InvalidPatchError(msg)

            typemodel_cls = self._get_typemodel_class(model_field.annotation)
            updated_value = value
            if typemodel_cls is not None and isinstance(updated_value, dict):
                updated_value = typemodel_cls(**updated_value)

            if updated_value is None and model_field.is_required():
                msg = f'Cannot set required field to null: {name}'
                raise InvalidPatchError(msg)

            setattr(target, name, updated_value)

        if self.on_before_save:
            await self.on_before_save(target)
        await target.asave()
        return target

    async def _resolve_nested_fields(
        self,
        fields: dict[str, object],
        context: ContextResolution,
        update_query: str,
    ) -> 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) or isinstance(value, Model):
                continue

            if isinstance(value, dict) and '$ref' in value:
                continue

            resolution = resolutions.get(field_name)
            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

            if isinstance(value, str) and value and self._is_field_mentioned(update_query, field_name):
                entity_cls = self._get_entity_class(entity_type)
                if entity_cls is None:
                    continue

                results = await self._search_entity(entity_cls, [value])
                if results:
                    resolved[field_name] = build_reference_payload(
                        entity_type,
                        str(results[0].object_id),
                    )

        return resolved

    def _filter_unmentioned_fk_fields(
        self,
        fields: dict[str, object],
        update_query: str,
    ) -> dict[str, object]:
        model_fields = self.model_class.model_fields
        filtered = dict(fields)

        for name in fields:
            field = model_fields.get(name)
            if field is None:
                continue
            model_cls = self._get_model_class(field.annotation)
            if model_cls is None:
                continue
            if self._has_related_update_intent(update_query, name, model_cls):
                filtered.pop(name, None)
                self.logger.debug('NLQ update dropped fk field due to related intent: %s', name)
                continue
            if self._is_field_mentioned(update_query, name):
                continue
            filtered.pop(name, None)
            self.logger.debug('NLQ update dropped fk field: %s', name)

        return filtered

    def _is_field_mentioned(self, update_query: str, field_name: str) -> bool:
        query = update_query.lower()
        if field_name.lower() in query:
            return True

        field = self.model_class.model_fields.get(field_name)
        if field is not None:
            title = field.title
            if isinstance(title, str) and title.lower() in query:
                return True

            model_cls = self._get_model_class(field.annotation)
            if model_cls is not None and model_cls.__name__.lower() in query:
                return True

        return False

    async def _normalize_related_update_query(
        self,
        update_query: str,
        fk_field_name: str,
        model_cls: type[Model],
    ) -> str:
        fk_field = self.model_class.model_fields.get(fk_field_name)
        fk_title = ''
        if fk_field is not None and isinstance(fk_field.title, str):
            fk_title = fk_field.title

        prompt = get_prompt('nl_query_update_related_normalize', custom_path=self.prompt_path)
        full_prompt = prompt.render_text(
            query=update_query,
            fk_field_name=fk_field_name,
            fk_field_title=fk_title,
            related_model=model_cls.__name__,
        )

        response_format = ResponseFormat.PLAIN_TEXT
        if ResponseFormat.PLAIN_TEXT not in self.llm.supported_formats:
            response_format = (
                ResponseFormat.JSON_OBJECT
                if ResponseFormat.JSON_OBJECT in self.llm.supported_formats
                else ResponseFormat.PLAIN_TEXT
            )

        try:
            response = await self.llm.ainvoke(
                full_prompt,
                response_format=response_format,
            )
        except Exception as exc:  # noqa: BLE001
            self.logger.warning('NLQ related normalize failed, using raw query: %s', exc)
            return update_query

        normalized = response.strip()
        if not normalized:
            return update_query

        return normalized


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

    Wraps `NLQueryUpdaterExecutor` and returns updated records as `Document` objects.

    Example:
        >>> llm = OpenAIModel()
        >>> llm.setup()
        >>> queryset = Product.objects.all()
        >>> updater = NLQueryUpdater(llm=llm, queryset=queryset)
        >>> analysis = await updater.analyze("update products set price to 99 where brand is ACME")
        >>> docs = await updater.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 = NLQueryUpdaterExecutor(
            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) -> UpdateAnalysis:
        return await self.executor.analyze(query, preview_limit=preview_limit)

    async def invoke(self, analysis: UpdateAnalysis) -> 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 | UpdateAnalysis, *, 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(),
        )
