"""LLM mocker implementation."""

import json
import logging
from typing import Any, Callable

from pydantic import BaseModel

from uipath.tracing import traced
from uipath.tracing._utils import _SpanUtils

from .._models._evaluation_set import EvaluationItem, LLMMockingStrategy
from .._models._mocks import ExampleCall
from .mocker import (
    Mocker,
    R,
    T,
    UiPathMockResponseGenerationError,
    UiPathNoMockFoundError,
)

PROMPT = """You are simulating a tool call for automated testing purposes of an Agent.
Your task is to provide a realistic response for the tool based on its schema, examples, and context.

SELECTED TOOL INFORMATION:
{toolInfo}

SELECTED TOOL SCHEMA:
{toolSchema}

SELECTED TOOL EXAMPLES:
{toolRunExamples}

CHOSEN TOOL INPUT:
{currentToolInput}

CURRENT AGENT RUN SO FAR:
{testRunHistory}

HERE IS SOME INFORMATION ABOUT THE AGENT: DO NOT USE THIS INFORMATION TO MAKE THE RESPONSE, BUT RATHER TO UNDERSTAND THE CONTEXT IN WHICH THE TOOL IS BEING USED.
{agentInfo}

TEST RUN PROCTOR INSTRUCTIONS:
You will need to simulate a real user's interaction with the tool.  This may require following some run specific instructions. If run instructions are provided, follow them exactly.
Here are the instructions for this run:
{testRunProctorInstructions}

Based on the above information, provide a realistic response for this tool call. Your response should:
1. Match the expected output format according to the tool schema
2. Be very consistent with how the tool has responded in previous examples. Do no omit fields or properties.
3. Always include the entire output regardless of token length.
3. Consider the context of the current test run and the agent being tested.  If the agent is acting on a property, make sure the output includes that property.

Respond ONLY with valid JSON that would be a realistic and complete tool response. Do not include any explanations or markdown.
"""

logger = logging.getLogger(__name__)


def pydantic_to_dict_safe(obj: Any) -> Any:
    """Serialize nested pydantic models to a dict."""
    if isinstance(obj, BaseModel):
        # Convert Pydantic model to dict recursively
        return obj.model_dump(mode="json")
    elif isinstance(obj, dict):
        # Recursively convert dict entries
        return {k: pydantic_to_dict_safe(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        # Recursively convert items in lists
        return [pydantic_to_dict_safe(item) for item in obj]
    return obj  # Return other types as is


class LLMMocker(Mocker):
    """LLM Based Mocker."""

    def __init__(self, evaluation_item: EvaluationItem):
        """LLM Mocker constructor."""
        self.evaluation_item = evaluation_item
        assert isinstance(self.evaluation_item.mocking_strategy, LLMMockingStrategy)

    @traced(name="__mocker__", recording=False)
    async def response(
        self, func: Callable[[T], R], params: dict[str, Any], *args: T, **kwargs
    ) -> R:
        """Respond with mocked response generated by an LLM."""
        assert isinstance(self.evaluation_item.mocking_strategy, LLMMockingStrategy)

        function_name = params.get("name") or func.__name__
        if function_name in [
            x.name for x in self.evaluation_item.mocking_strategy.tools_to_simulate
        ]:
            from uipath import UiPath
            from uipath._services.llm_gateway_service import _cleanup_schema

            from .mocks import (
                evaluation_context,
                execution_id_context,
                span_collector_context,
            )

            llm = UiPath().llm
            return_type: Any = func.__annotations__.get("return", None)
            if return_type is None:
                return_type = Any

            class OutputSchema(BaseModel):
                response: return_type

            response_format = {
                "type": "json_schema",
                "json_schema": {
                    "name": OutputSchema.__name__.lower(),
                    "strict": True,
                    "schema": _cleanup_schema(OutputSchema),
                },
            }
            try:
                # Safely pull examples from params.
                example_calls = params.get("example_calls", [])
                if not isinstance(example_calls, list):
                    example_calls = []
                example_calls = [
                    call for call in example_calls if isinstance(call, ExampleCall)
                ]

                test_run_history = "(empty)"
                eval_item = evaluation_context.get()
                span_collector = span_collector_context.get()
                execution_id = execution_id_context.get()
                if eval_item and span_collector and execution_id:
                    spans = span_collector.get_spans(execution_id)
                    test_run_history = _SpanUtils.spans_to_llm_context(spans)

                prompt_input: dict[str, Any] = {
                    "toolRunExamples": example_calls,
                    "testRunHistory": test_run_history,
                    "toolInfo": {
                        "name": function_name,
                        "description": params.get("description"),
                        "arguments": params.get(
                            "arguments"
                        ),  # arguments could be passed into tool
                        "settings": params.get(
                            "settings"
                        ),  # settings could be passed into tool
                        "inputSchema": params.get("input_schema"),
                    },
                    "toolSchema": params.get("input_schema"),
                    "currentToolInput": {
                        "args": args,
                        "kwargs": kwargs,
                    },
                    "agentInfo": {  # This is incomplete
                        # "agentName": self.evaluation_item.name, # to be obtained.
                        "actionName": self.evaluation_item.name,  # Not sure if this is right?
                        "userInput": self.evaluation_item.inputs,
                    },
                    "testRunProctorInstructions": self.evaluation_item.mocking_strategy.prompt,
                }
                prompt_input = {
                    k: json.dumps(pydantic_to_dict_safe(v))
                    for k, v in prompt_input.items()
                }
                model_parameters = self.evaluation_item.mocking_strategy.model
                completion_kwargs = (
                    model_parameters.model_dump(by_alias=False, exclude_none=True)
                    if model_parameters
                    else {}
                )
                response = await llm.chat_completions(
                    [
                        {
                            "role": "user",
                            "content": PROMPT.format(**prompt_input),
                        },
                    ],
                    response_format=response_format,
                    **completion_kwargs,
                )
                mocked_response = OutputSchema(
                    **json.loads(response.choices[0].message.content)
                )
                return mocked_response.model_dump(mode="json")["response"]
            except Exception as e:
                raise UiPathMockResponseGenerationError() from e
        else:
            raise UiPathNoMockFoundError(f"Method '{function_name}' is not simulated.")
