import datetime as dt
import json
import re
from dataclasses import dataclass
from typing import Any, AsyncIterator

import litellm

from freeact import tracing
from freeact.model.base import CodeActModel, CodeActModelResponse, CodeActModelTurn, CodeActModelUsage
from freeact.model.prompt import (
    CODE_TAG_SYSTEM_TEMPLATE,
    EXECUTION_ERROR_TEMPLATE,
    EXECUTION_OUTPUT_TEMPLATE,
    TOOL_USE_SYSTEM_TEMPLATE,
)
from freeact.model.tools import (
    CODE_EDITOR_TOOL,
    CODE_EXECUTOR_TOOL,
    beta_flag,
    code_editor_tool,
    code_editor_tool_use_default,
    code_executor_tool,
    code_executor_tool_use_default,
    sanitize_tool_name,
    tool_name,
)


@dataclass
class ToolUse:
    id: str
    name: str
    input: dict[str, Any]


@dataclass
class LiteLLMResponse(CodeActModelResponse):
    tool_use: ToolUse | None = None

    @property
    def tool_use_id(self) -> str | None:
        return self.tool_use.id if self.tool_use else None

    @property
    def tool_use_name(self) -> str | None:
        return self.tool_use.name if self.tool_use else None

    @property
    def code(self) -> str | None:
        if self.is_error:
            return None
        elif self.tool_use_name == tool_name(CODE_EXECUTOR_TOOL):
            return self.tool_use.input["code"]  # type: ignore
        elif self.tool_use_name == tool_name(CODE_EDITOR_TOOL):
            return f"print(file_editor(**{self.tool_use.input}))"  # type: ignore
        else:
            return self.code_block()

    def code_block(self) -> str | None:
        pattern = r"<code-action>\s*```python\s*([\s\S]*?)\s*```\s*</code-action>"
        match = re.search(pattern, self.text)
        return match.group(1) if match else None


class LiteLLMTurn(CodeActModelTurn):
    def __init__(self, iter: AsyncIterator[str | LiteLLMResponse], span_name: str, span_input: dict[str, Any]):
        self._iter = iter
        self._response: LiteLLMResponse | None = None
        self._span_name = span_name
        self._span_input = span_input

    async def response(self) -> LiteLLMResponse:
        if self._response is None:
            async for _ in self.stream():
                pass
        return self._response  # type: ignore

    async def stream(self) -> AsyncIterator[str]:
        async with tracing.span(
            self._span_name,
            self._span_input,
        ) as span:
            await span.update(start_time=dt.datetime.now())
            async for elem in self._iter:
                match elem:
                    case str():
                        yield elem
                    case LiteLLMResponse() as msg:
                        self._response = msg
                        await span.update(
                            output={
                                "text": msg.text,
                                "is_error": msg.is_error,
                                "code": msg.code,
                                "tool_use": msg.tool_use,
                            },
                        )


class LiteCodeActModel(CodeActModel):
    """A [LiteLLM](https://github.com/BerriAI/litellm)-based code action model.

    Code actions are generated differently depending on the `use_executor_tool` argument:

    - `use_executor_tool=False`: Code actions are included directly into the model's response text,
      enclosed in `<code-action> ... </code-action>` tags. Uses the
      [`CODE_TAG_SYSTEM_TEMPLATE`][freeact.model.CODE_TAG_SYSTEM_TEMPLATE] by default.

    - `use_executor_tool=True`: Code actions are generated by calling an internal `execute_ipython_cell`
      tool. Uses the [`TOOL_USE_SYSTEM_TEMPLATE`][freeact.model.TOOL_USE_SYSTEM_TEMPLATE] by default.

    - `use_executor_tool=None`: A sensible default is chosen based on the model name and provider.
      Currently, a tool use approach is used for Anthropic and OpenAI models, a code tag approach
      is used for all other models.

    A custom system template can be provided with the `system_template` constructor argument.
    Its semantics should match the `use_executor_tool` argument value.

    Models created with `use_editor_tool=True` are also able to create and edit files. This
    allows them to store and edit code actions on disk (= long-term memory). Stored code actions can
    be loaded as *custom skills* via [`get_sources`][freeact.environment.CodeProvider.get_sources].
    If `use_editor_tool` is `None`, a sensible default is chosen based on the model name and provider.
    Currently, Anthropic and OpenAI models are configured to use the editor tool.

    Args:
        model_name: A model name supported by [LiteLLM](https://github.com/BerriAI/litellm).
        skill_sources: Source code of Python modules offered to the model as skills. They
            are utilized by generated code actions, if useful for the task. Skill sources are
            usually loaded and formatted with [`get_sources`][freeact.environment.CodeProvider.get_sources].
        system_template: A system template that guides the model to generate code actions.
            Must define a `{python_modules}` placeholder for `skill_sources`.
        execution_output_template: A prompt template for formatting successful code execution
            output. Must define an `{execution_feedback}` placeholder.
        execution_error_template: A prompt template for formatting code execution errors.
            Must define an `{execution_feedback}` placeholder.
        use_executor_tool: Whether to use a tool-based approach for generating code actions (`True`)
            or a code tag based approach (`False`). If `None`, a sensible default is chosen based on
            `model_name`.
        use_editor_tool: Whether to use a file editor tool for creating and editing code action modules
            on disk. If `None`, a sensible default is chosen based on `model_name`.
        **kwargs: Default [chat completion](https://docs.litellm.ai/docs/completion) `kwargs`
            for [`request`][freeact.model.LiteCodeActModel.request] and
            [`feedback`][freeact.model.LiteCodeActModel.feedback] calls.
            These are merged with `request` and `feedback` specific completion `kwargs`
            where the latter have higher precedence in case of conflicting keys. The following
            `kwargs` are set internally and must not be set here: `stream`, `stream_options`,
            `messages`, and `tools`.
    """

    def __init__(
        self,
        model_name: str,
        skill_sources: str | None = None,
        system_template: str | None = None,
        execution_output_template: str | None = None,
        execution_error_template: str | None = None,
        use_executor_tool: bool | None = None,
        use_editor_tool: bool | None = None,
        **kwargs,
    ):
        self.model_name = model_name
        self.completion_kwargs = kwargs

        if "max_tokens" not in self.completion_kwargs:
            self.completion_kwargs["max_tokens"] = 8192

        if use_executor_tool is None:
            use_executor_tool = code_executor_tool_use_default(model_name, self.provider_name)

        if use_editor_tool is None:
            use_editor_tool = code_editor_tool_use_default(model_name, self.provider_name)

        if execution_output_template is None:
            execution_output_template = EXECUTION_OUTPUT_TEMPLATE

        if execution_error_template is None:
            execution_error_template = EXECUTION_ERROR_TEMPLATE

        if system_template is None:
            system_template = TOOL_USE_SYSTEM_TEMPLATE if use_executor_tool else CODE_TAG_SYSTEM_TEMPLATE

        self.execution_output_template = execution_output_template
        self.execution_error_template = execution_error_template

        system_instruction = system_template.format(python_modules=skill_sources or "")

        if self.provider_name == "anthropic":
            if self.completion_kwargs.pop("prompt_caching", True):
                system_instruction = [  # type: ignore
                    {
                        "type": "text",
                        "text": system_instruction,
                        "cache_control": {
                            "type": "ephemeral",
                        },
                    }
                ]

        self.history: list[dict[str, Any]] = [{"role": "system", "content": system_instruction}]
        self.tools: list[dict[str, Any]] = []

        if use_executor_tool:
            self.tools.append(code_executor_tool(model_name))

        if use_editor_tool:
            self.tools.append(code_editor_tool(model_name))

            if flag := beta_flag(model_name):
                self.completion_kwargs["extra_headers"] = flag

    @property
    def tool_names(self) -> list[str]:
        """The names of the tools configured for this model."""
        return [tool_name(tool) for tool in self.tools]

    @property
    def provider_name(self) -> str:
        """The name of the model's provider."""
        if "/" in self.model_name:
            provider, *_ = self.model_name.split("/")
            endpoint = self.completion_kwargs.get("api_base") or self.completion_kwargs.get("base_url")
            if provider == "openai" and endpoint:
                return "openai-compat"
            return provider
        else:
            return "openai"

    def request(
        self,
        user_query: str,
        **kwargs,
    ) -> CodeActModelTurn:
        """Initiates an interaction with this model from a user query.

        Args:
            user_query: The user query (a question, instruction, etc.)
            **kwargs: [Chat completion](https://docs.litellm.ai/docs/completion) arguments.
                These are merged with the model's default completion `kwargs`. Default
                completion `kwargs` have lower precedence in case of conflicting keys.

        Returns:
            CodeActModelTurn: An object for retrieving the model's response.
        """
        user_message = {"role": "user", "content": user_query}

        span_name = "Model request"
        span_input = {"user_query": user_query, **kwargs}

        return LiteLLMTurn(self._stream(user_message, **kwargs), span_name, span_input)

    def feedback(
        self,
        feedback: str,
        is_error: bool,
        tool_use_id: str | None,
        tool_use_name: str | None,
        **kwargs,
    ) -> CodeActModelTurn:
        """Initiates an interaction with this model from code execution feedback,
        allowing the model to refine or correct previous responses, or returning
        a final response to the user. A `feedback` call must follow a previous
        `request` or `feedback` call.

        Args:
            feedback: The feedback text from code execution or other actions.
            is_error: Whether the `feedback` text contains error information.
            **kwargs: [Chat completion](https://docs.litellm.ai/docs/completion) arguments.
                These are merged with the model's default completion `kwargs`. Default
                completion `kwargs` have lower precedence in case of conflicting keys.

        Returns:
            CodeActModelTurn: An object for retrieving the model's response.
        """
        if tool_use_name == tool_name(CODE_EXECUTOR_TOOL) or tool_use_name is None:
            template = self.execution_error_template if is_error else self.execution_output_template
            content = template.format(execution_feedback=feedback)
        else:  # skip application of execution feedback templates for results of other tools
            content = feedback

        if tool_use_id is None:
            feedback_message = {
                "role": "user",
                "content": content,
            }
        else:
            feedback_message = {
                "role": "tool",
                "tool_call_id": tool_use_id,
                "content": content,
            }

        span_name = "Model feedback"
        span_input = {
            "feedback": feedback,
            "is_error": is_error,
            "tool_use_id": tool_use_id,
            "tool_use_name": tool_use_name,
            **kwargs,
        }

        return LiteLLMTurn(self._stream(feedback_message, **kwargs), span_name, span_input)

    async def _stream(self, input_message: dict[str, Any], **kwargs) -> AsyncIterator[str | LiteLLMResponse]:
        span = tracing.get_active_span()
        if span.trace_id and span.span_id:
            metadata = {
                "existing_trace_id": span.trace_id,
                "parent_observation_id": span.span_id,
            }
        else:
            metadata = {}

        messages = self.history + [input_message]
        response = LiteLLMResponse(text="", is_error=False)

        result_stream = await litellm.acompletion(
            model=self.model_name,
            messages=messages,
            stream=True,
            stream_options={"include_usage": True},
            tools=self.tools or None,
            **(self.completion_kwargs | kwargs),
            metadata=metadata,
        )

        chunks = []
        chunk_deltas = []
        think = False

        async for chunk in result_stream:
            chunks.append(chunk)
            chunk_delta = chunk.choices[0].delta
            chunk_deltas.append(chunk_delta)

            if hasattr(chunk_delta, "reasoning_content") and chunk_delta.reasoning_content:
                if not think:
                    think = True
                    yield "<think>\n"
                yield chunk_delta.reasoning_content

            if hasattr(chunk_delta, "content") and chunk_delta.content:
                if think:
                    think = False
                    yield "\n</think>\n\n"
                yield chunk_delta.content

        result = litellm.stream_chunk_builder(chunks, messages=messages)
        result_message = result.choices[0].message

        # litellm.stream_chunk_builder() does not include the thinking blocks
        # emitted by Anthropic models so we need to accumulate them here.
        if thinking_block := self._extract_thinking_block(chunk_deltas):
            result_message.thinking_blocks = [thinking_block]

        # message to be added to history ...
        assistant_message: dict[str, Any] = {
            "role": "assistant",
        }

        if hasattr(result_message, "reasoning_content") and result_message.reasoning_content:
            response.text += f"<think>\n{result_message.reasoning_content}\n</think>\n\n"

        if content := result_message.content:
            response.text += content

        if content := self._extract_content(result_message):
            assistant_message["content"] = content

        if result_message.tool_calls:
            tool_call = result_message.tool_calls[0]
            tool_name = sanitize_tool_name(tool_call.function.name)

            if tool_name not in self.tool_names:
                allowed_tool_names = ", ".join(self.tool_names)
                response.is_error = True
                response.text = f"Invalid tool name: {tool_name}\nAllowed tool names are: {allowed_tool_names}"

            try:
                tool_input = json.loads(tool_call.function.arguments)
            except json.JSONDecodeError:
                tool_input = {}
                response.is_error = True
                response.text = f"Could not decode tool input: {tool_call.function.arguments}"

            response.tool_use = ToolUse(
                id=tool_call.id,
                name=tool_name,
                input=tool_input,
            )

            assistant_message["tool_calls"] = [
                {
                    "id": tool_call.id,
                    "type": "function",
                    "function": {
                        "name": tool_name,
                        "arguments": tool_call.function.arguments,
                    },
                }
            ]

        response.usage = self._extract_usage(result)

        self.history.append(input_message)
        self.history.append(assistant_message)

        yield response

    @staticmethod
    def _extract_thinking_block(chunks_deltas) -> dict[str, Any] | None:
        # Accumulates thinking block chunks generated by Anthropic models

        # -----------------------------------------------------------------------------------------------------
        #  FIXME: handle redacted_thinking blocks when https://github.com/BerriAI/litellm/pull/10329 is merged
        # -----------------------------------------------------------------------------------------------------

        thinking = ""
        signature = None

        for chunk_delta in chunks_deltas:
            if not hasattr(chunk_delta, "thinking_blocks"):
                continue

            for block in chunk_delta.thinking_blocks:
                if block.get("type") == "thinking":
                    if thinking_delta := block.get("thinking"):
                        thinking += thinking_delta

                    if signature_delta := block.get("signature"):
                        signature = signature_delta

        return {"type": "thinking", "thinking": thinking, "signature": signature} if thinking else None

    @staticmethod
    def _extract_content(message: litellm.Message):
        # Extracts message content that preserves thinking blocks generated by Anthropic models

        if hasattr(message, "thinking_blocks") and message.thinking_blocks:
            # Anthropic-specific content structure
            output = message.thinking_blocks
            if content := message.content:
                output.append(
                    {
                        "type": "text",
                        "text": content,
                    },
                )
            return output
        else:
            # Content is an optional string
            return message.content

    @staticmethod
    def _extract_usage(response: litellm.ModelResponse) -> CodeActModelUsage:
        # Extracts usage and cost information from a model response

        usage = response.usage
        result = CodeActModelUsage(
            input_tokens=usage.prompt_tokens,
            output_tokens=usage.completion_tokens,
            total_tokens=usage.total_tokens,
        )

        if cache_write_tokens := usage.get("cache_creation_input_tokens"):
            result.cache_write_tokens = cache_write_tokens
        if cache_read_tokens := usage.get("cache_read_input_tokens"):
            result.cache_read_tokens = cache_read_tokens

        if prompt_tokens_details := usage.get("prompt_tokens_details"):
            if cached_tokens := prompt_tokens_details.cached_tokens:
                result.cache_read_tokens = cached_tokens

        if completion_tokens_details := usage.get("completion_tokens_details"):
            if thinking_tokens := completion_tokens_details.reasoning_tokens:
                result.thinking_tokens = thinking_tokens

        try:
            result.cost = litellm.completion_cost(completion_response=response)
        except Exception:
            pass

        return result
