from __future__ import annotations

import logging
from typing import Any

from amsdal_models.classes.model import Model
from pydantic import ValidationError

from amsdal_ml.mcp_server.mcp_auth import authorize_class
from amsdal_ml.mcp_server.mcp_auth import authorize_object
from amsdal_ml.mcp_server.mcp_auth import with_amsdal_context
from amsdal_ml.mcp_server.server_model_explorer import mcp_server
from amsdal_ml.mcp_server.server_model_explorer.services.crud_utils import get_action
from amsdal_ml.mcp_server.server_model_explorer.services.crud_utils import get_api_queryset
from amsdal_ml.mcp_server.server_model_explorer.services.dependencies import _get_llm
from amsdal_ml.mcp_server.server_model_explorer.services.dependencies import _get_model_registry
from amsdal_ml.mcp_server.server_model_explorer.services.types import CrudAnalysisResponse
from amsdal_ml.mcp_server.server_model_explorer.services.types import CrudOperation
from amsdal_ml.mcp_server.server_model_explorer.services.types import CrudResponse
from amsdal_ml.ml_retrievers.query_creator import CreateAnalysis
from amsdal_ml.ml_retrievers.query_creator import NLQueryCreatorExecutor
from amsdal_ml.ml_retrievers.query_deleter import DeleteAnalysis
from amsdal_ml.ml_retrievers.query_deleter import NLQueryDeleterExecutor
from amsdal_ml.ml_retrievers.query_retriever import NLQueryExecutor
from amsdal_ml.ml_retrievers.query_retriever import RetrieveAnalysis
from amsdal_ml.ml_retrievers.query_updater import NLQueryUpdaterExecutor
from amsdal_ml.ml_retrievers.query_updater import UpdateAnalysis
from amsdal_ml.utils.query_utils import serialize_and_clean_record

logger = logging.getLogger(__name__)


@mcp_server.tool(
    name='perform_crud_operation',
    description=(
        'Performs CRUD operations on AMSDAL data models using natural language queries with a two-step confirm flow. '
        'ALL operations (read, create, update, delete) require confirmation. '
        'Step 1: call with confirm=false to get analysis and requires_confirmation=true. '
        'Step 2: after user approval, call again with confirm=true and pass the analysis payload to execute. '
        'On the first call you must show preview results and ask the user to confirm the intent. '
        'Repeat the tool call with confirm=true only after the user confirms. '
        'The analysis contains preview records only and does NOT change the database. '
        'Changes happen only when confirm=true. '
        'IMPORTANT: The natural language query processors automatically handle nested entity '
        'creation, updates, and references. '
        'Do NOT create separate operations for nested entities - include all related data in a single query. '
        'For example, when creating an order record that references a customer, specify the '
        'customer details in the same query. '
        'The system will resolve existing entities or create new ones as needed. '
        'If the user asks to show/verify a specific object by id '
        '(for example: "before I confirm, show book by id ..."), '
        'use get_model_field_values_by_ids instead of operation=read. '
        'Format final answers using Markdown tables.'
    ),
)
@with_amsdal_context
async def perform_crud_operation(
    operation: CrudOperation,
    model_name: str,
    query: str,
    *,
    confirm: bool = False,
    preview_limit: int = 5,
    limit: int | None = 30,
    analysis: dict[str, Any] | None = None,
) -> CrudResponse | CrudAnalysisResponse | str:
    registry = _get_model_registry()
    model_info = registry.get(model_name)

    if not model_info:
        return f"Error: Model '{model_name}' not found."
    model_cls = model_info.model

    llm = _get_llm()

    try:
        action = get_action(operation)
        await authorize_class(model_cls, action)
    except Exception as e:
        return f'Authorization error: {e!s}'

    if operation == CrudOperation.READ:
        try:

            async def _authorize_read(obj):
                await authorize_object(obj, action)

            queryset = get_api_queryset(model_cls)

            nlq_executor = NLQueryExecutor(
                llm=llm,
                queryset=queryset,
                on_object_loaded=_authorize_read,
            )

            return await _run_crud_confirm_flow(
                operation=operation,
                model_name=model_name,
                executor=nlq_executor,
                analysis_model=RetrieveAnalysis,
                query=query,
                confirm=confirm,
                analysis_payload=analysis,
                analyze_kwargs={'query': query, 'preview_limit': preview_limit, 'limit': limit},
            )
        except ValidationError as exc:
            logger.warning(
                'Invalid analysis payload: op=%s model=%s confirm=%s error=%s',
                operation.value,
                model_name,
                confirm,
                exc,
            )
            return f'Error: invalid analysis payload: {exc!s}'
        except Exception as e:
            logger.exception(
                'Error executing CRUD operation: op=%s model=%s confirm=%s query=%r',
                operation.value,
                model_name,
                confirm,
                query,
            )
            return f'Error executing query: {e!s}'

    if operation == CrudOperation.CREATE:
        try:

            async def _authorize_create(obj):
                await authorize_object(obj, action)

            queryset = get_api_queryset(model_cls)

            create_executor = NLQueryCreatorExecutor(
                llm=llm,
                queryset=queryset,
                on_before_save=_authorize_create,
            )

            return await _run_crud_confirm_flow(
                operation=operation,
                model_name=model_name,
                executor=create_executor,
                analysis_model=CreateAnalysis,
                query=query,
                confirm=confirm,
                analysis_payload=analysis,
                analyze_kwargs={'query': query, 'preview_limit': preview_limit},
            )
        except ValidationError as exc:
            logger.warning(
                'Invalid analysis payload: op=%s model=%s confirm=%s error=%s',
                operation.value,
                model_name,
                confirm,
                exc,
            )
            return f'Error: invalid analysis payload: {exc!s}'
        except Exception as e:
            logger.exception(
                'Error executing CRUD operation: op=%s model=%s confirm=%s query=%r',
                operation.value,
                model_name,
                confirm,
                query,
            )
            return f'Error executing query: {e!s}'

    if operation == CrudOperation.UPDATE:
        try:

            async def _authorize_update(obj):
                await authorize_object(obj, action)

            queryset = get_api_queryset(model_cls)

            update_executor = NLQueryUpdaterExecutor(
                llm=llm,
                queryset=queryset,
                on_before_save=_authorize_update,
            )
            return await _run_crud_confirm_flow(
                operation=operation,
                model_name=model_name,
                executor=update_executor,
                analysis_model=UpdateAnalysis,
                query=query,
                confirm=confirm,
                analysis_payload=analysis,
                analyze_kwargs={'query': query, 'preview_limit': preview_limit},
            )
        except ValidationError as exc:
            logger.warning(
                'Invalid analysis payload: op=%s model=%s confirm=%s error=%s',
                operation.value,
                model_name,
                confirm,
                exc,
            )
            return f'Error: invalid analysis payload: {exc!s}'
        except Exception as e:
            logger.exception(
                'Error executing CRUD operation: op=%s model=%s confirm=%s query=%r',
                operation.value,
                model_name,
                confirm,
                query,
            )
            return f'Error executing query: {e!s}'

    if operation == CrudOperation.DELETE:
        try:

            async def _authorize_delete(obj):
                await authorize_object(obj, action)

            queryset = get_api_queryset(model_cls)

            delete_executor = NLQueryDeleterExecutor(
                llm=llm,
                queryset=queryset,
                on_before_delete=_authorize_delete,
            )

            return await _run_crud_confirm_flow(
                operation=operation,
                model_name=model_name,
                executor=delete_executor,
                analysis_model=DeleteAnalysis,
                query=query,
                confirm=confirm,
                analysis_payload=analysis,
                analyze_kwargs={'query': query, 'preview_limit': preview_limit},
            )
        except ValidationError as exc:
            logger.warning(
                'Invalid analysis payload: op=%s model=%s confirm=%s error=%s',
                operation.value,
                model_name,
                confirm,
                exc,
            )
            return f'Error: invalid analysis payload: {exc!s}'
        except Exception as e:
            logger.exception(
                'Error executing CRUD operation: op=%s model=%s confirm=%s query=%r',
                operation.value,
                model_name,
                confirm,
                query,
            )
            return f'Error executing query: {e!s}'

    return f"Operation '{operation.value}' is not yet implemented."


async def _run_crud_confirm_flow(
    *,
    operation: CrudOperation,
    model_name: str,
    executor: Any,
    analysis_model: Any,
    query: str,
    confirm: bool,
    analysis_payload: dict[str, Any] | None,
    analyze_kwargs: dict[str, Any],
) -> CrudResponse | CrudAnalysisResponse:
    """Shared two-step confirm flow for NLQ CRUD operations."""

    logger.debug('CRUD confirm flow: op=%s model=%s confirm=%s query=%r', operation.value, model_name, confirm, query)

    if confirm:
        if analysis_payload is None:
            msg = "The 'analysis' payload is required when confirm=True."
            raise ValueError(msg)

        analysis_obj = analysis_model.model_validate(analysis_payload)
        results: list[Model] = await executor.execute(analysis_obj)
        cleaned_results = [await serialize_and_clean_record(r) for r in results]

        return CrudResponse(
            operation=operation.value,
            model=model_name,
            data=cleaned_results,
        )

    analysis_obj = await executor.analyze(**analyze_kwargs)

    return CrudAnalysisResponse(
        operation=operation.value,
        model=model_name,
        analysis=analysis_obj.model_dump(),
        requires_confirmation=True,
    )
