import logging
import warnings
from datetime import datetime, timezone
from typing import Any, AsyncGenerator, Dict, List, Set, Tuple, Union, cast

from arthur_common.models.llm_model_providers import (
    LLMResponseFormat,
    OpenAIMessage,
    OpenAIMessageType,
    ToolChoiceEnum,
)
from fastapi.responses import StreamingResponse
from jinja2 import meta
from jinja2.sandbox import SandboxedEnvironment
from litellm import (
    CustomStreamWrapper,
    Message,
    stream_chunk_builder,
)
from litellm.types.utils import ModelResponse, ModelResponseStream

from clients.llm.llm_client import LLMClient, LLMModelResponse
from schemas.agentic_prompt_schemas import AgenticPrompt
from schemas.enums import SSEEventType
from schemas.request_schemas import CompletionRequest, PromptCompletionRequest
from schemas.response_schemas import AgenticPromptRunResponse
from utils.sse_events import format_sse, format_sse_error

# litellm's Anthropic streaming handler sets Choices instead of StreamingChoices on chunks,
# causing pydantic serialization warnings on model_dump_json() and completion_cost().
# This is a litellm bug — suppress it process-wide to avoid noise in logs.
warnings.filterwarnings(
    "ignore",
    message="Pydantic serializer warnings",
    category=UserWarning,
)

logger = logging.getLogger(__name__)


class ChatCompletionService:
    """Service for running a chat completion"""

    def __init__(self) -> None:
        # autoescape=False because Jinja automatically HTML-escapes items by default. Ex:
        #   if text = "{{ name }}" and variables = {"name": "<Bob>"}
        #   autoescape=False "{{ name }}" -> "<Bob>"
        #   autoescape=True "{{ name }}" -> "&lt;Bob&gt;"
        self._jinja_env = SandboxedEnvironment(autoescape=False)

    # ============================================================================
    # Variable Validation and Replacement
    # ============================================================================

    def find_undeclared_variables_in_text(
        self,
        text: str,
    ) -> Set[str]:
        template_ast = self._jinja_env.parse(text)
        return meta.find_undeclared_variables(template_ast)

    def _find_missing_variables_in_variable_map(
        self,
        variable_map: Dict[str, str],
        text: str,
    ) -> List[str]:
        undeclared_vars = self.find_undeclared_variables_in_text(text)
        return [v for v in undeclared_vars if v not in variable_map]

    def find_missing_variables_in_messages(
        self,
        variable_map: Dict[str, str],
        messages: List[OpenAIMessage],
    ) -> Set[str]:
        missing_vars = set()
        for message in messages:
            if message.content is not None:
                if isinstance(message.content, str):
                    missing_vars.update(
                        self._find_missing_variables_in_variable_map(
                            variable_map,
                            message.content,
                        ),
                    )
                elif isinstance(message.content, list):
                    for item in message.content:
                        if item.type == OpenAIMessageType.TEXT.value and item.text:
                            missing_vars.update(
                                self._find_missing_variables_in_variable_map(
                                    variable_map,
                                    item.text,
                                ),
                            )

            for tool_call in message.tool_calls or []:
                missing_vars.update(
                    self._find_missing_variables_in_variable_map(
                        variable_map,
                        tool_call.function.arguments,
                    ),
                )

        return missing_vars

    def _replace_variables_in_text(
        self,
        variable_map: Dict[str, str],
        text: str,
    ) -> str:
        template = self._jinja_env.from_string(text)
        return template.render(**variable_map)

    def replace_variables(
        self,
        variable_map: Dict[str, str],
        messages: List[OpenAIMessage],
    ) -> list[OpenAIMessage]:
        for message in messages:
            if message.content is not None:
                if isinstance(message.content, str):
                    message.content = self._replace_variables_in_text(
                        variable_map,
                        message.content,
                    )
                elif isinstance(message.content, list):
                    for item in message.content:
                        if item.type == OpenAIMessageType.TEXT.value and item.text:
                            item.text = self._replace_variables_in_text(
                                variable_map,
                                item.text,
                            )

            for tool_call in message.tool_calls or []:
                tool_call.function.arguments = self._replace_variables_in_text(
                    variable_map,
                    tool_call.function.arguments,
                )

        return messages

    @staticmethod
    def to_prompt_and_request(
        unsaved_prompt: CompletionRequest,
    ) -> Tuple[AgenticPrompt, PromptCompletionRequest]:
        """
        Convert an unsaved run request into its corresponding AgenticPrompt and PromptCompletionRequest
        """
        prompt = AgenticPrompt(
            name="test_unsaved_prompt",
            created_at=datetime.now(timezone.utc),
            **unsaved_prompt.model_dump(exclude={"completion_request"}),
        )
        return prompt, unsaved_prompt.completion_request

    # ============================================================================
    # Completion Helpers and Methods
    # ============================================================================

    @staticmethod
    def _extract_token_counts(
        response: Any,
    ) -> tuple[None | int, None | int, None | int]:
        """
        Extract token counts from a response object's usage attribute.

        Args:
            response: Response object (ModelResponse, TextCompletionResponse, etc.) that may have a usage attribute

        Returns:
            Tuple of (input_tokens, output_tokens, total_tokens), all may be None
        """
        input_tokens = None
        output_tokens = None
        total_tokens = None

        if hasattr(response, "usage") and response.usage:
            usage = response.usage
            input_tokens = getattr(usage, "prompt_tokens", None)
            output_tokens = getattr(usage, "completion_tokens", None)
            total_tokens = getattr(usage, "total_tokens", None)

        return input_tokens, output_tokens, total_tokens

    def _get_completion_params(
        self,
        prompt: AgenticPrompt,
        completion_request: PromptCompletionRequest = PromptCompletionRequest(),
    ) -> Tuple[str, Dict[str, Any]]:
        model = (
            prompt.model_name
            if prompt.model_name.startswith(prompt.model_provider + "/")
            else prompt.model_provider + "/" + prompt.model_name
        )

        completion_params: dict[str, Any] = {
            "messages": [
                message.model_dump(exclude_none=True) for message in prompt.messages
            ],
        }

        if prompt.tools:
            completion_params["tools"] = [
                tool.model_dump(exclude_none=True) for tool in prompt.tools
            ]

        # flatten config params to pass into litellm properly
        if prompt.config:
            if prompt.config.timeout:
                completion_params["timeout"] = prompt.config.timeout
            if prompt.config.temperature:
                completion_params["temperature"] = prompt.config.temperature
            if prompt.config.top_p:
                completion_params["top_p"] = prompt.config.top_p
            if prompt.config.max_tokens:
                completion_params["max_tokens"] = prompt.config.max_tokens
            if prompt.config.stop:
                completion_params["stop"] = prompt.config.stop
            if prompt.config.presence_penalty:
                completion_params["presence_penalty"] = prompt.config.presence_penalty
            if prompt.config.frequency_penalty:
                completion_params["frequency_penalty"] = prompt.config.frequency_penalty
            if prompt.config.seed:
                completion_params["seed"] = prompt.config.seed
            if prompt.config.logprobs:
                completion_params["logprobs"] = prompt.config.logprobs
            if prompt.config.top_logprobs:
                completion_params["top_logprobs"] = prompt.config.top_logprobs
            if prompt.config.logit_bias:
                completion_params["logit_bias"] = [
                    logit_bias.model_dump(exclude_none=True)
                    for logit_bias in prompt.config.logit_bias
                ]
            if prompt.config.max_completion_tokens:
                completion_params["max_completion_tokens"] = (
                    prompt.config.max_completion_tokens
                )
            if prompt.config.reasoning_effort:
                completion_params["reasoning_effort"] = (
                    prompt.config.reasoning_effort.value
                )
            if prompt.config.thinking:
                completion_params["thinking"] = dict(prompt.config.thinking)
            if prompt.config.tool_choice:
                if isinstance(prompt.config.tool_choice, ToolChoiceEnum):
                    completion_params["tool_choice"] = prompt.config.tool_choice.value
                else:
                    completion_params["tool_choice"] = (
                        prompt.config.tool_choice.model_dump(exclude_none=True)
                    )
            if prompt.config.response_format:
                # either get the json_object/schema response format or pass in the pydantic model
                response_format = prompt.config.response_format
                if isinstance(response_format, LLMResponseFormat):
                    completion_params["response_format"] = response_format.model_dump(
                        exclude_none=True,
                    )
                else:
                    completion_params["response_format"] = response_format
            if prompt.config.stream_options:
                completion_params["stream_options"] = (
                    prompt.config.stream_options.model_dump(exclude_none=True)
                )

        # validate all variables are passed in to the prompt if strict mode is enabled
        if completion_request.strict == True:
            missing_vars = self.find_missing_variables_in_messages(
                completion_request._variable_map,
                prompt.messages,
            )
            if missing_vars:
                raise ValueError(
                    f"Missing values for the following variables: {', '.join(sorted(missing_vars))}",
                )

        # replace variables in messages
        completion_messages = prompt.messages
        if completion_request.variables:
            completion_messages = self.replace_variables(
                completion_request._variable_map,
                completion_messages,
            )

        completion_params["messages"] = [
            message.model_dump(exclude_none=True) for message in completion_messages
        ]

        if completion_request.stream is not None:
            completion_params["stream"] = completion_request.stream

        # not allowed to have tool_choice if no tools are provided
        if prompt.tools is None:
            completion_params.pop("tool_choice", None)

        return model, completion_params

    def run_chat_completion_raw_response(
        self,
        prompt: AgenticPrompt,
        llm_client: LLMClient,
        completion_request: PromptCompletionRequest = PromptCompletionRequest(),
    ) -> LLMModelResponse:
        model, completion_params = self._get_completion_params(
            prompt,
            completion_request,
        )
        return llm_client.completion(model=model, **completion_params)

    def run_chat_completion(
        self,
        prompt: AgenticPrompt,
        llm_client: LLMClient,
        completion_request: PromptCompletionRequest = PromptCompletionRequest(),
    ) -> AgenticPromptRunResponse:
        if prompt.has_been_deleted():
            raise ValueError(
                f"Cannot run chat completion for this prompt because it was deleted on: {prompt.deleted_at}",
            )

        llm_model_response = self.run_chat_completion_raw_response(
            prompt,
            llm_client,
            completion_request,
        )
        if not isinstance(llm_model_response.response, ModelResponse):
            raise ValueError("Response is not a ModelResponse")
        if hasattr(llm_model_response.response.choices[0], "message"):
            msg: Message = cast(
                Message,
                getattr(llm_model_response.response.choices[0], "message"),
            )
            # Extract token counts from usage if available
            input_tokens, output_tokens, total_tokens = self._extract_token_counts(
                llm_model_response.response,
            )

            return AgenticPromptRunResponse(
                content=msg.content,
                tool_calls=msg.tool_calls,
                cost=llm_model_response.cost or "",
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                total_tokens=total_tokens,
            )
        else:
            raise ValueError("No message from model")

    async def stream_chat_completion(
        self,
        prompt: AgenticPrompt,
        llm_client: LLMClient,
        completion_request: PromptCompletionRequest = PromptCompletionRequest(),
    ) -> AsyncGenerator[str, None]:
        try:
            if prompt.has_been_deleted():
                raise ValueError(
                    f"Cannot stream chat completion for this prompt because it was deleted on: {prompt.deleted_at}",
                )

            model, completion_params = self._get_completion_params(
                prompt,
                completion_request,
            )
            response = await llm_client.acompletion(model=model, **completion_params)

            collected_chunks: list[
                ModelResponse | ModelResponseStream | CustomStreamWrapper
            ] = []
            if isinstance(response, CustomStreamWrapper):
                async for chunk in response:
                    collected_chunks.append(chunk)
                    yield f"event: chunk\ndata: {chunk.model_dump_json()}\n\n"
            elif isinstance(response, ModelResponse):
                collected_chunks.append(response)
            else:
                raise NotImplementedError(
                    f"Unsupported response type: {type(response)}",
                )

            complete_response = stream_chunk_builder(
                collected_chunks,
                messages=completion_params.get("messages", []),
            )

            cost = llm_client.calculate_cost(complete_response)

            if not complete_response:
                yield format_sse_error("No response from model", wrap=False)
            elif not complete_response.choices:
                yield format_sse_error("No choices from model", wrap=False)
            elif hasattr(complete_response.choices[0], "message"):
                msg: Message = cast(
                    Message,
                    getattr(complete_response.choices[0], "message"),
                )
                # Extract token counts from usage if available
                input_tokens, output_tokens, total_tokens = self._extract_token_counts(
                    complete_response,
                )

                data = AgenticPromptRunResponse(
                    content=msg.content,
                    tool_calls=msg.tool_calls,
                    cost=cost,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    total_tokens=total_tokens,
                ).model_dump_json()

                yield format_sse(SSEEventType.FINAL_RESPONSE, data)
        except Exception as e:
            yield format_sse_error(str(e), wrap=False)

    async def execute_prompt_completion(
        self,
        llm_client: LLMClient,
        prompt: AgenticPrompt,
        completion_request: PromptCompletionRequest,
    ) -> Union[AgenticPromptRunResponse, StreamingResponse]:
        """Helper to execute prompt completion with or without streaming"""
        if completion_request.stream is None or completion_request.stream == False:
            return self.run_chat_completion(prompt, llm_client, completion_request)

        return StreamingResponse(
            self.stream_chat_completion(prompt, llm_client, completion_request),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
        )
