from __future__ import annotations

import json
import logging
from typing import Any

from amsdal_models.classes.model import Model
from amsdal_server.apps.common.permissions.enums import Action
from amsdal_server.apps.transactions.services.transaction_execution_api import TransactionExecutionApi
from amsdal_utils.models.utils.reference_builders import build_reference
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_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.models import ErrorArgAnalysis
from amsdal_ml.mcp_server.server_model_explorer.services.models import ExistingRefArgAnalysis
from amsdal_ml.mcp_server.server_model_explorer.services.models import ExplicitIdArgAnalysis
from amsdal_ml.mcp_server.server_model_explorer.services.models import LiteralArgAnalysis
from amsdal_ml.mcp_server.server_model_explorer.services.models import NewObjectArgAnalysis
from amsdal_ml.mcp_server.server_model_explorer.services.models import PlannedArguments
from amsdal_ml.mcp_server.server_model_explorer.services.models import TransactionArgAnalysis
from amsdal_ml.mcp_server.server_model_explorer.services.models import TransactionRunAnalysis
from amsdal_ml.mcp_server.server_model_explorer.services.models import TransactionSpec
from amsdal_ml.mcp_server.server_model_explorer.services.transaction_utils import collect_transaction_catalog
from amsdal_ml.mcp_server.server_model_explorer.services.transaction_utils import is_system_file_model
from amsdal_ml.mcp_server.server_model_explorer.services.transaction_utils import resolve_transaction_name
from amsdal_ml.mcp_server.server_model_explorer.services.types import TransactionAnalysisResponse
from amsdal_ml.mcp_server.server_model_explorer.services.types import TransactionArgumentMode
from amsdal_ml.mcp_server.server_model_explorer.services.types import TransactionExecutionResponse
from amsdal_ml.mcp_server.server_model_explorer.services.types import TransactionParameterKind
from amsdal_ml.mcp_server.server_model_explorer.services.types import TransactionParameterView
from amsdal_ml.mcp_server.server_model_explorer.services.types import TransactionScope
from amsdal_ml.ml_models.utils import ResponseFormat
from amsdal_ml.ml_retrievers.adapters import get_retriever_adapter
from amsdal_ml.ml_retrievers.exceptions import MultipleRecordsError
from amsdal_ml.ml_retrievers.query_creator import NLQueryCreatorExecutor
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 resolve_reference_resource
from amsdal_ml.utils.response_format import select_response_format

logger = logging.getLogger(__name__)


@mcp_server.tool(
    name='perform_transaction_operation',
    description=(
        'Executes AMSDAL transactions (custom database operations/functions) from natural language descriptions. '
        'MANDATORY two-step process: '
        'Step 1: Call with confirm=false to analyze the request and preview what will be executed. '
        'Step 2: After user confirms, call with confirm=true and pass the returned analysis payload '
        '(analysis is required when confirm=true). '
        'CRITICAL: Pass the user\'s EXACT raw input as the "query" parameter - '
        'DO NOT modify, rephrase, summarize, or interpret it. '
        "Copy the user's message verbatim, including typos, abbreviations, and informal language. "
        'Examples of what NOT to do: '
        '- User says: "add book xyz qty 5" -> DO NOT send "Add book with title xyz and quantity 5" '
        '- User says: "call AddToCart book(existing): title abc author def qty 123" -> '
        'DO NOT send "book: {title: abc, author: def}, quantity: 123" '
        'Always send the raw user input exactly as typed. '
        'If the user adds clarifications or corrections in later messages, preserve prior intent '
        'and append the new context '
        'instead of replacing or truncating the original request. '
        'Context-aware is required, paraphrasing is forbidden: combine user messages faithfully, do not rewrite them. '
        'This raw query is fed directly to AI for argument planning, so precision matters.'
    ),
)
@with_amsdal_context
async def perform_transaction_operation(
    transaction_name: str,
    query: str,
    *,
    confirm: bool = False,
    analysis: dict[str, Any] | None = None,
    preview_limit: int = 5,
) -> TransactionAnalysisResponse | TransactionExecutionResponse | str:
    try:
        catalog = await collect_transaction_catalog(scope=TransactionScope.ALL)
    except Exception as exc:
        logger.exception('Error collecting transaction catalog')
        return f'Error loading transactions: {exc!s}'

    try:
        resolved_name = resolve_transaction_name(catalog, transaction_name)
    except ValueError as exc:
        return f'Error resolving transaction: {exc!s}'

    transaction = catalog[resolved_name]

    try:
        if confirm:
            if analysis is None:
                return "Error: The 'analysis' payload is required when confirm=True."

            run_analysis = TransactionRunAnalysis.model_validate(analysis)
            if not run_analysis.is_executable:
                return 'Error: transaction arguments are ambiguous. Ask user to clarify and retry confirm=false.'

            args = await _materialize_transaction_args(
                transaction=transaction,
                run_analysis=run_analysis,
            )

            result = await TransactionExecutionApi.execute_transaction(
                transaction_name=resolved_name,
                args=args,
            )

            return {'transaction_name': resolved_name, 'result': result}

        run_analysis = await _analyze_transaction_request(
            transaction=transaction,
            query=query,
            preview_limit=preview_limit,
        )

        return {
            'transaction_name': resolved_name,
            'analysis': run_analysis.model_dump(),
            'requires_confirmation': True,
        }
    except ValidationError as exc:
        logger.warning('Invalid transaction analysis payload: transaction=%s error=%s', resolved_name, exc)
        return f'Error: invalid analysis payload: {exc!s}'
    except Exception as exc:
        logger.exception('Error executing transaction operation: transaction=%s confirm=%s', resolved_name, confirm)
        return f'Error executing transaction: {exc!s}'


async def _analyze_transaction_request(
    *,
    transaction: TransactionSpec,
    query: str,
    preview_limit: int,
) -> TransactionRunAnalysis:
    """
    Analyze user query and produce executable argument plan per parameter.

    This stage validates argument modes against parameter kinds and prepares
    resolution payloads for model arguments (existing reference or new object)
    without executing the transaction itself.

    Args:
        transaction: Transaction metadata with typed parameter specs.
        query: Raw user request used for argument planning and fallback search.
        preview_limit: Max preview records used in NLQ analyzers.

    Returns:
        TransactionRunAnalysis: Per-argument analysis and overall executability.
    """

    llm = _get_llm()
    planned = await _plan_arguments(transaction=transaction, query=query)
    planned_by_name = {item.name: item for item in planned.items}

    arg_analyses: list[TransactionArgAnalysis] = []
    is_executable = True

    for param in transaction.parameters:
        arg = planned_by_name.get(param.name)
        kind = param.kind
        model_name = param.model_name

        if arg is None:
            arg_analyses.append(
                ErrorArgAnalysis(
                    name=param.name,
                    kind=kind,
                    model_name=model_name,
                    error='Missing argument plan from LLM',
                )
            )
            is_executable = False
            continue

        if kind in {TransactionParameterKind.PRIMITIVE, TransactionParameterKind.DICT}:
            if arg.mode != TransactionArgumentMode.LITERAL:
                arg_analyses.append(
                    ErrorArgAnalysis(
                        name=param.name,
                        kind=kind,
                        model_name=model_name,
                        error=f'Expected literal mode for parameter {param.name}',
                    )
                )
                is_executable = False
            else:
                arg_analyses.append(
                    LiteralArgAnalysis(
                        name=param.name,
                        kind=kind,
                        model_name=model_name,
                        value=arg.literal_value,
                    )
                )
            continue

        if kind != TransactionParameterKind.MODEL or model_name is None:
            arg_analyses.append(
                ErrorArgAnalysis(
                    name=param.name,
                    kind=kind,
                    model_name=model_name,
                    error='Unknown parameter kind. Please provide explicit argument.',
                )
            )
            is_executable = False
            continue

        model_info = _get_model_registry().get(model_name)
        if model_info is None:
            arg_analyses.append(
                ErrorArgAnalysis(
                    name=param.name,
                    kind=kind,
                    model_name=model_name,
                    error=f'Model {model_name} not found in registry',
                )
            )
            is_executable = False
            continue

        model_cls = model_info.model

        if arg.mode == TransactionArgumentMode.LITERAL:
            if param.nullable and arg.literal_value is None:
                arg_analyses.append(
                    LiteralArgAnalysis(
                        name=param.name,
                        kind=kind,
                        model_name=model_name,
                        value=None,
                    )
                )
            else:
                arg_analyses.append(
                    ErrorArgAnalysis(
                        name=param.name,
                        kind=kind,
                        model_name=model_name,
                        error=f'Unsupported mode {arg.mode} for model parameter',
                    )
                )
                is_executable = False
            continue

        if arg.mode == TransactionArgumentMode.NEW_OBJECT and is_system_file_model(model_cls):
            arg_analyses.append(
                ErrorArgAnalysis(
                    name=param.name,
                    kind=kind,
                    model_name=model_name,
                    error='new_object mode is not supported for File model. Use explicit_id or existing_ref.',
                )
            )
            is_executable = False
            continue

        if arg.mode == TransactionArgumentMode.EXPLICIT_ID:
            if not arg.explicit_object_id:
                arg_analyses.append(
                    ErrorArgAnalysis(
                        name=param.name,
                        kind=kind,
                        model_name=model_name,
                        error='explicit_id mode requires explicit_object_id',
                    )
                )
                is_executable = False
            else:
                arg_analyses.append(
                    ExplicitIdArgAnalysis(
                        name=param.name,
                        kind=kind,
                        model_name=model_name,
                        value=arg.explicit_object_id,
                    )
                )
            continue

        if arg.mode == TransactionArgumentMode.EXISTING_REF:
            await authorize_class(model_cls, Action.READ)

            async def _authorize_read(obj: Model) -> None:
                await authorize_object(obj, Action.READ)

            search_query = arg.query or query
            queryset = get_api_queryset(model_cls)

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

            retrieval_analysis = await executor.analyze(
                search_query,
                preview_limit=preview_limit,
                limit=2,
            )

            retrieval_error: str | None = None
            if retrieval_analysis.total != 1:
                retrieval_error = (
                    f'Expected exactly 1 {model_name} for existing_ref, got {retrieval_analysis.total}. '
                    'Need clarification.'
                )
                is_executable = False
            arg_analyses.append(
                ExistingRefArgAnalysis(
                    name=param.name,
                    kind=kind,
                    model_name=model_name,
                    query=search_query,
                    existing_analysis=retrieval_analysis,
                    error=retrieval_error,
                )
            )
            continue

        if arg.mode == TransactionArgumentMode.NEW_OBJECT:
            await authorize_class(model_cls, Action.CREATE)

            async def _authorize_create(obj: Model) -> None:
                await authorize_object(obj, Action.CREATE)

            create_query = arg.query or query
            queryset = get_api_queryset(model_cls)

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

            create_analysis = await creator.analyze(
                create_query,
                preview_limit=preview_limit,
            )

            create_error: str | None = None
            if create_analysis.total != 1:
                create_error = (
                    f'Expected exactly 1 payload for new_object {model_name}, got {create_analysis.total}. '
                    'Need clarification.'
                )
                is_executable = False
            arg_analyses.append(
                NewObjectArgAnalysis(
                    name=param.name,
                    kind=kind,
                    model_name=model_name,
                    query=create_query,
                    create_analysis=create_analysis,
                    error=create_error,
                )
            )
            continue

        arg_analyses.append(
            ErrorArgAnalysis(
                name=param.name,
                kind=kind,
                model_name=model_name,
                error=f'Unsupported mode {arg.mode} for model parameter',
            )
        )
        is_executable = False

    summary = 'ready' if is_executable else 'needs_clarification'
    return TransactionRunAnalysis(
        transaction_name=transaction.name,
        args=arg_analyses,
        is_executable=is_executable,
        summary=summary,
    )


async def _materialize_transaction_args(
    *,
    transaction: TransactionSpec,
    run_analysis: TransactionRunAnalysis,
) -> dict[str, Any]:
    """
    Convert validated analysis into concrete transaction execution arguments.

    Primitive/dict values are passed as-is. Model values are resolved to
    references via explicit id, existing object retrieval, or object creation,
    based on the planned argument mode.

    Args:
        transaction: Target transaction specification.
        run_analysis: Confirmed argument analysis payload.

    Returns:
        dict[str, Any]: Execution-ready argument mapping for Transaction API.

    Raises:
        ValueError: If analysis is incomplete, invalid, or ambiguous.
    """
    llm = _get_llm()
    args: dict[str, Any] = {}
    analysis_by_name = {item.name: item for item in run_analysis.args}

    for param in transaction.parameters:
        name = param.name
        item = analysis_by_name.get(name)
        if item is None:
            msg = f'Missing analysis for parameter {name}'
            raise ValueError(msg)
        if item.error:
            msg = f'Argument {name} is not executable: {item.error}'
            raise ValueError(msg)

        if item.kind in {TransactionParameterKind.PRIMITIVE, TransactionParameterKind.DICT}:
            if not isinstance(item, LiteralArgAnalysis):
                msg = f'Unsupported state for {name}: expected literal argument'
                raise ValueError(msg)
            args[name] = item.value
            continue

        if item.kind != TransactionParameterKind.MODEL:
            msg = f'Unsupported parameter kind for {name}: {item.kind}'
            raise ValueError(msg)
        if item.model_name is None:
            msg = f'Model name is missing for {name}: {item.kind}'
            raise ValueError(msg)

        model_info = _get_model_registry().get(item.model_name)
        if model_info is None:
            msg = f'Model {item.model_name} not found in registry'
            raise ValueError(msg)
        model_cls = model_info.model

        if isinstance(item, LiteralArgAnalysis):
            if param.nullable and item.value is None:
                args[name] = None
                continue
            msg = f'Unsupported mode for {name}: {item.mode}'
            raise ValueError(msg)

        if isinstance(item, ExplicitIdArgAnalysis):
            args[name] = build_reference(
                class_name=item.model_name,
                object_id=item.value,
                resource=resolve_reference_resource(item.model_name),
            )
            continue

        if isinstance(item, ExistingRefArgAnalysis):
            if item.existing_analysis is None:
                msg = f'existing_ref analysis is missing for {name}'
                raise ValueError(msg)

            async def _authorize_read(obj: Model) -> None:
                await authorize_object(obj, Action.READ)

            retrieval_analysis = item.existing_analysis
            queryset = get_api_queryset(model_cls)

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

            try:
                model_obj = await executor.execute_one(retrieval_analysis)
            except MultipleRecordsError as exc:
                preview = json.dumps(exc.records[:5], ensure_ascii=False)
                msg = f'Expected 1 resolved object for {name}, got {exc.total}. Preview: {preview}'
                raise ValueError(msg) from exc

            args[name] = build_reference(
                class_name=item.model_name,
                object_id=model_obj.object_id,
                resource=resolve_reference_resource(item.model_name),
            )
            continue

        if isinstance(item, NewObjectArgAnalysis):
            if item.create_analysis is None:
                msg = f'new_object analysis is missing for {name}'
                raise ValueError(msg)

            async def _authorize_create(obj: Model) -> None:
                await authorize_object(obj, Action.CREATE)

            create_analysis = item.create_analysis
            queryset = get_api_queryset(model_cls)

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

            try:
                created_obj = await creator.execute_one(create_analysis)
            except MultipleRecordsError as exc:
                preview = json.dumps(exc.records[:5], ensure_ascii=False)
                msg = f'Expected 1 created object for {name}, got {exc.total}. Preview: {preview}'
                raise ValueError(msg) from exc

            args[name] = build_reference(
                class_name=item.model_name,
                object_id=created_obj.object_id,
                resource=resolve_reference_resource(item.model_name),
            )
            continue

        msg = f'Unsupported mode for {name}: {item.mode}'
        raise ValueError(msg)

    return args


async def _plan_arguments(
    *,
    transaction: TransactionSpec,
    query: str,
) -> PlannedArguments:
    """
    Use LLM to map raw user query to structured transaction argument plan.

    Builds planner prompt from transaction parameter metadata and parses model
    output into ``PlannedArguments`` using either JSON schema or plain-text JSON
    cleanup flow, depending on model capabilities.

    Args:
        transaction: Transaction view used to describe expected arguments.
        query: Raw user query passed to the planner prompt.

    Returns:
        PlannedArguments: Structured argument plan produced by the model.

    Raises:
        ValueError: If model output is not valid JSON or fails validation.
    """
    llm = _get_llm()
    prompt = get_prompt('transaction_argument_planner')
    parameters_json = json.dumps(
        [
            TransactionParameterView(
                name=parameter.name,
                annotation=parameter.annotation,
                kind=parameter.kind,
                model_name=parameter.model_name,
                nullable=parameter.nullable,
            )
            for parameter in transaction.parameters
        ],
        ensure_ascii=False,
        indent=2,
    )
    full_prompt = prompt.render_text(
        transaction_name=transaction.name,
        parameters_json=parameters_json,
        query=query,
    )

    response_format = select_response_format(llm.supported_formats)
    schema: dict[str, Any] | None = None
    if response_format == ResponseFormat.JSON_SCHEMA:
        adapter = get_retriever_adapter(llm)
        schema = adapter.get_response_schema(PlannedArguments.model_json_schema())

    response = await llm.ainvoke(
        full_prompt,
        response_format=response_format,
        schema=schema,
    )

    raw = response.strip()
    if response_format == ResponseFormat.PLAIN_TEXT:
        raw = clean_json_response(raw)

    try:
        data = json.loads(raw)
        return PlannedArguments.model_validate(data)
    except (json.JSONDecodeError, ValidationError) as exc:
        msg = f'Failed to parse transaction argument plan: {exc}'
        raise ValueError(msg) from exc
