from __future__ import annotations

import json
import logging
from collections import defaultdict
from collections.abc import AsyncIterator, Callable, Iterator
from typing import Any

from any_llm import acompletion, completion
from any_llm.types.completion import (
    ChatCompletion,
    ChatCompletionChunk,
    ChatCompletionMessage,
    ReasoningEffort,
)

from django_anyllm.prompts import build_system_prompt
from django_anyllm.registry import ModelRegistry, default_registry
from django_anyllm.settings import get_extra_tools
from django_anyllm.stream_types import StreamEvent, StreamEventType
from django_anyllm.tools import get_tools

logger = logging.getLogger(__name__)


def _infer_provider(model: str) -> str | None:
    """Infer a provider from a bare model name."""
    model_lower = model.lower()
    if model_lower.startswith(("gpt-", "o1", "o3", "text-")):
        return "openai"
    if model_lower.startswith("claude-"):
        return "anthropic"
    if model_lower.startswith("kimi-"):
        return "moonshot"
    if model_lower.startswith("deepseek-"):
        return "deepseek"
    if model_lower.startswith("llama-"):
        return "meta"
    if model_lower.startswith("mistral-"):
        return "mistral"
    return None


def _resolve_model(model: str, provider: str | None) -> str:
    """Return a model identifier suitable for any-llm.

    If *provider* is given and *model* has no prefix, returns
    ``provider:model``.  If *provider* is blank but the model name
    matches a known provider, infers and prepends the provider.
    Otherwise returns *model* unchanged.
    """
    if ":" in model or "/" in model:
        return model

    if provider:
        return f"{provider}:{model}"

    inferred = _infer_provider(model)
    if inferred:
        return f"{inferred}:{model}"

    return model


def _strip_unsupported_file_blocks(
    messages: list[dict[str, Any]], provider: str | None
) -> list[dict[str, Any]]:
    """Remove file content blocks not supported by the target provider.

    Anthropic only accepts application/pdf as a document media type.  Any
    other file blocks in the history would cause a 400, so they are dropped
    before the request is sent.  Other providers are left unchanged.
    """
    if provider != "anthropic":
        return messages

    result = []
    for msg in messages:
        content = msg.get("content")
        if not isinstance(content, list):
            result.append(msg)
            continue
        filtered = []
        for block in content:
            if block.get("type") == "file":
                file_data: str = block.get("file", {}).get("file_data", "")
                mime = file_data[5:].split(";")[0] if file_data.startswith("data:") else ""
                if mime != "application/pdf":
                    continue
            filtered.append(block)
        result.append({**msg, "content": filtered} if len(filtered) != len(content) else msg)
    return result


def _strip_files_from_history(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Remove file/image content blocks from all but the last user message.

    Re-sending base64-encoded files for every turn of the conversation is
    extremely token-heavy.  The model already saw those files in earlier turns,
    so stripping them from older messages saves the bulk of the input tokens
    without losing any conversational context.
    """
    last_user_idx = next(
        (i for i in range(len(messages) - 1, -1, -1) if messages[i].get("role") == "user"),
        None,
    )
    result = []
    for i, msg in enumerate(messages):
        if i == last_user_idx or not isinstance(msg.get("content"), list):
            result.append(msg)
            continue
        filtered = [
            block for block in msg["content"]
            if block.get("type") not in ("image_url", "file")
        ]
        result.append({**msg, "content": filtered} if len(filtered) != len(msg["content"]) else msg)
    return result


def _collect_streaming_tool_calls(
    chunks: list[ChatCompletionChunk],
) -> list[dict[str, Any]]:
    """Reconstruct complete tool calls from streaming chunks.

    OpenAI-style streaming sends tool call fragments across multiple chunks.
    We accumulate them by index to form complete tool call objects.
    """
    calls: dict[int, dict[str, Any]] = defaultdict(
        lambda: {"id": None, "function": {"name": "", "arguments": ""}}
    )

    for chunk in chunks:
        delta = chunk.choices[0].delta if chunk.choices else None
        if not delta or not delta.tool_calls:
            continue
        for tc in delta.tool_calls:
            idx = tc.index
            if tc.id:
                calls[idx]["id"] = tc.id
            if tc.function:
                if tc.function.name:
                    calls[idx]["function"]["name"] = tc.function.name
                if tc.function.arguments:
                    calls[idx]["function"]["arguments"] += tc.function.arguments

    return list(calls.values())


class ModelAgent:
    """Orchestrate LLM conversations with registered Django models.

    Example::

        agent = ModelAgent(model="gpt-4", provider="openai")
        answer = agent.chat("List all authors")

        for event in agent.chat("List all authors", stream=True):
            print(event.text or "", end="")
    """

    def __init__(
        self,
        model: str,
        provider: str | None = None,
        *,
        registry: ModelRegistry | None = None,
        system_prompt: str | None = None,
        max_tool_rounds: int = 5,
        api_key: str | None = None,
        api_base: str | None = None,
        temperature: float | None = None,
        **kwargs: Any,
    ) -> None:
        """Initialise the agent.

        Args:
            model: Model identifier (e.g. ``gpt-4``).
            provider: Provider name (e.g. ``openai``). If omitted, inferred from
                the model string.
            registry: Model registry to use. Defaults to the global registry.
            system_prompt: Override the auto-generated system prompt.
            max_tool_rounds: Maximum number of back-and-forth tool call rounds.
            api_key: Provider API key.
            api_base: Provider base URL.
            temperature: Sampling temperature.
            **kwargs: Additional arguments forwarded to ``any_llm.completion``.
        """
        self.model = model
        self.provider = provider
        self.registry = registry or default_registry
        self.system_prompt = system_prompt or build_system_prompt(self.registry)
        self.max_tool_rounds = max_tool_rounds
        self.api_key = api_key
        self.api_base = api_base
        self.temperature = temperature
        self.extra_kwargs = kwargs

        self._tools = get_tools(self.registry) + get_extra_tools()
        self._tool_map: dict[str, Callable[..., Any]] = {
            t.__name__: t for t in self._tools
        }
        self._history: list[dict[str, Any]] = []

    def _build_messages(self, user_message: str) -> list[dict[str, Any]]:
        return [
            {"role": "system", "content": self.system_prompt},
            *self._history,
            {"role": "user", "content": user_message},
        ]

    def _execute_tool(self, name: str, args: dict[str, Any]) -> Any:
        """Execute a single tool call and return its result."""
        func = self._tool_map.get(name)
        if func is None:
            return f"Error: tool '{name}' not found."
        try:
            return func(**args)
        except Exception as exc:
            logger.exception("Tool execution failed: %s", name)
            return f"Error executing {name}: {exc}"

    def _run_tool_calls(self, tool_calls: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Execute a batch of tool calls and return tool result messages."""
        results: list[dict[str, Any]] = []
        for tc in tool_calls:
            name = tc["function"]["name"]
            try:
                args = json.loads(tc["function"]["arguments"])
            except json.JSONDecodeError as exc:
                args = {}
                result = f"Error: invalid JSON arguments: {exc}"
            else:
                result = self._execute_tool(name, args)
            results.append(
                {
                    "role": "tool",
                    "tool_call_id": tc.get("id", ""),
                    "name": name,
                    "content": json.dumps(result)
                    if not isinstance(result, str)
                    else result,
                }
            )
        return results

    def _extract_assistant_message(
        self, response: ChatCompletion
    ) -> ChatCompletionMessage:
        return response.choices[0].message

    def _message_to_dict(self, message: ChatCompletionMessage) -> dict[str, Any]:
        """Convert a ChatCompletionMessage to a plain dict for the next request."""
        data: dict[str, Any] = {
            "role": message.role,
            "content": message.content or "",
        }
        if message.tool_calls:
            data["tool_calls"] = [
                {
                    "id": tc.id,
                    "type": tc.type,
                    "function": {
                        "name": tc.function.name,
                        "arguments": tc.function.arguments,
                    },
                }
                for tc in message.tool_calls
            ]
        return data

    def chat(
        self,
        user_message: str,
        *,
        stream: bool = False,
        reasoning_effort: ReasoningEffort | None = None,
    ) -> str | Iterator[StreamEvent]:
        """Send a user message and receive the assistant's response.

        Args:
            user_message: The user's message.
            stream: If ``True``, returns a generator that yields
                :class:`~django_anyllm.stream_types.StreamEvent` objects.
            reasoning_effort: Reasoning level for supported models.

        Returns:
            The final assistant text (non-streaming) or a generator
            of stream events.
        """
        if stream:
            return self._chat_stream(user_message, reasoning_effort=reasoning_effort)
        return self._chat_sync(user_message, reasoning_effort=reasoning_effort)

    def _chat_sync(
        self, user_message: str, *, reasoning_effort: ReasoningEffort | None = None
    ) -> str:
        """Non-streaming chat with automatic tool-call loop."""
        messages = self._build_messages(user_message)

        for _round in range(self.max_tool_rounds):
            response: ChatCompletion = completion(
                model=_resolve_model(self.model, self.provider),
                messages=messages,
                tools=self._tools,
                reasoning_effort=reasoning_effort,
                api_key=self.api_key,
                api_base=self.api_base,
                temperature=self.temperature,
                **self.extra_kwargs,
            )

            message = self._extract_assistant_message(response)
            messages.append(self._message_to_dict(message))

            if not message.tool_calls:
                return message.content or ""

            tool_results = self._run_tool_calls(
                [
                    {
                        "id": tc.id,
                        "function": {
                            "name": tc.function.name,
                            "arguments": tc.function.arguments,
                        },
                    }
                    for tc in message.tool_calls
                ]
            )
            messages.extend(tool_results)

        # Safety limit reached.
        return message.content or ""

    def _chat_stream(
        self, user_message: str, *, reasoning_effort: ReasoningEffort | None = None
    ) -> Iterator[StreamEvent]:
        """Streaming chat with automatic tool-call loop.

        Yields content and reasoning chunks in real time, then tool call
        and tool result events when the model requests tools.
        """
        messages = self._build_messages(user_message)

        for _round in range(self.max_tool_rounds):
            chunks: list[ChatCompletionChunk] = []
            stream = completion(
                model=_resolve_model(self.model, self.provider),
                messages=messages,
                tools=self._tools,
                stream=True,
                reasoning_effort=reasoning_effort,
                api_key=self.api_key,
                api_base=self.api_base,
                temperature=self.temperature,
                **self.extra_kwargs,
            )

            for chunk in stream:
                chunks.append(chunk)
                choice = chunk.choices[0] if chunk.choices else None
                if not choice:
                    continue

                delta = choice.delta
                if delta.content:
                    yield StreamEvent(type=StreamEventType.CONTENT, text=delta.content)
                if delta.reasoning and delta.reasoning.content:
                    yield StreamEvent(
                        type=StreamEventType.REASONING,
                        text=delta.reasoning.content,
                    )

            # Reconstruct the assistant message from chunks.
            tool_calls = _collect_streaming_tool_calls(chunks)
            content_parts: list[str] = []
            for chunk in chunks:
                delta = chunk.choices[0].delta if chunk.choices else None
                if delta and delta.content:
                    content_parts.append(delta.content)
            content = "".join(content_parts)

            assistant_msg: dict[str, Any] = {
                "role": "assistant",
                "content": content,
            }
            if tool_calls:
                assistant_msg["tool_calls"] = tool_calls
            messages.append(assistant_msg)

            if not tool_calls:
                break

            # Execute tools and yield events.
            for tc in tool_calls:
                name = tc["function"]["name"]
                try:
                    args = json.loads(tc["function"]["arguments"])
                except json.JSONDecodeError:
                    args = {}
                yield StreamEvent(
                    type=StreamEventType.TOOL_CALL,
                    tool_name=name,
                    tool_args=args,
                )

            tool_results = self._run_tool_calls(tool_calls)
            for tc, result_msg in zip(tool_calls, tool_results):
                yield StreamEvent(
                    type=StreamEventType.TOOL_RESULT,
                    tool_name=tc["function"]["name"],
                    tool_result=result_msg["content"],
                )

            messages.extend(tool_results)

        yield StreamEvent(type=StreamEventType.DONE)

    # ------------------------------------------------------------------
    # Async variants
    # ------------------------------------------------------------------

    async def achat(
        self,
        user_message: str,
        *,
        stream: bool = False,
        reasoning_effort: ReasoningEffort | None = None,
    ) -> str | AsyncIterator[StreamEvent]:
        """Async version of :meth:`chat`."""
        if stream:
            return self._achat_stream(user_message, reasoning_effort=reasoning_effort)
        return await self._achat_sync(user_message, reasoning_effort=reasoning_effort)

    async def _achat_sync(
        self, user_message: str, *, reasoning_effort: ReasoningEffort | None = None
    ) -> str:
        """Non-streaming async chat with automatic tool-call loop."""
        messages = self._build_messages(user_message)

        for _round in range(self.max_tool_rounds):
            response: ChatCompletion = await acompletion(
                model=_resolve_model(self.model, self.provider),
                messages=messages,
                tools=self._tools,
                reasoning_effort=reasoning_effort,
                api_key=self.api_key,
                api_base=self.api_base,
                temperature=self.temperature,
                **self.extra_kwargs,
            )

            message = self._extract_assistant_message(response)
            messages.append(self._message_to_dict(message))

            if not message.tool_calls:
                return message.content or ""

            tool_results = self._run_tool_calls(
                [
                    {
                        "id": tc.id,
                        "function": {
                            "name": tc.function.name,
                            "arguments": tc.function.arguments,
                        },
                    }
                    for tc in message.tool_calls
                ]
            )
            messages.extend(tool_results)

        return message.content or ""

    async def _achat_stream(
        self, user_message: str, *, reasoning_effort: ReasoningEffort | None = None
    ) -> AsyncIterator[StreamEvent]:
        """Streaming async chat with automatic tool-call loop."""
        messages = self._build_messages(user_message)

        for _round in range(self.max_tool_rounds):
            chunks: list[ChatCompletionChunk] = []
            stream = await acompletion(
                model=_resolve_model(self.model, self.provider),
                messages=messages,
                tools=self._tools,
                stream=True,
                reasoning_effort=reasoning_effort,
                api_key=self.api_key,
                api_base=self.api_base,
                temperature=self.temperature,
                **self.extra_kwargs,
            )

            async for chunk in stream:
                chunks.append(chunk)
                choice = chunk.choices[0] if chunk.choices else None
                if not choice:
                    continue

                delta = choice.delta
                if delta.content:
                    yield StreamEvent(type=StreamEventType.CONTENT, text=delta.content)
                if delta.reasoning and delta.reasoning.content:
                    yield StreamEvent(
                        type=StreamEventType.REASONING,
                        text=delta.reasoning.content,
                    )

            tool_calls = _collect_streaming_tool_calls(chunks)
            content_parts = [
                chunk.choices[0].delta.content
                for chunk in chunks
                if chunk.choices
                and chunk.choices[0].delta
                and chunk.choices[0].delta.content
            ]
            content = "".join(content_parts)

            assistant_msg: dict[str, Any] = {
                "role": "assistant",
                "content": content,
            }
            if tool_calls:
                assistant_msg["tool_calls"] = tool_calls
            messages.append(assistant_msg)

            if not tool_calls:
                break

            for tc in tool_calls:
                name = tc["function"]["name"]
                try:
                    args = json.loads(tc["function"]["arguments"])
                except json.JSONDecodeError:
                    args = {}
                yield StreamEvent(
                    type=StreamEventType.TOOL_CALL,
                    tool_name=name,
                    tool_args=args,
                )

            tool_results = self._run_tool_calls(tool_calls)
            for tc, result_msg in zip(tool_calls, tool_results):
                yield StreamEvent(
                    type=StreamEventType.TOOL_RESULT,
                    tool_name=tc["function"]["name"],
                    tool_result=result_msg["content"],
                )

            messages.extend(tool_results)

        yield StreamEvent(type=StreamEventType.DONE)


# ---------------------------------------------------------------------------
# Persistent agent
# ---------------------------------------------------------------------------


class PersistentModelAgent(ModelAgent):
    """A ModelAgent that persists messages to the database.

    Example::

        from django_anyllm.models import Chat

        chat = Chat.objects.create(model="gpt-4", provider="openai")
        agent = PersistentModelAgent(chat=chat)
        answer = agent.chat("List all authors")
    """

    def __init__(self, chat, **kwargs: Any) -> None:
        """Initialise the persistent agent.

        Args:
            chat: A ``Chat`` model instance (or subclass).
            **kwargs: Forwarded to :class:`ModelAgent`.
        """
        self._chat = chat
        system_prompt = chat.system_prompt or build_system_prompt(
            kwargs.get("registry") or default_registry
        )
        super().__init__(
            model=chat.model,
            provider=chat.provider or None,
            system_prompt=system_prompt,
            **kwargs,
        )

    @staticmethod
    def _db_message_to_llm(message) -> dict[str, Any]:
        """Convert a database Message to an any-llm-sdk message dict."""
        content: str | list[dict[str, Any]] = message.content or ""

        attached_files = list(message.files.all())
        if attached_files:
            import os

            parts: list[dict[str, Any]] = []
            file_meta = ", ".join(
                f"{os.path.basename(f.file.name)} (file_id={f.pk})" for f in attached_files
            )
            text = f"{content}\n[Attached files: {file_meta}]" if content else f"[Attached files: {file_meta}]"
            parts.append({"type": "text", "text": text})
            for f in attached_files:
                part = f.to_image_content() or f.to_file_content()
                if part:
                    parts.append(part)
            content = parts if parts else ""

        msg: dict[str, Any] = {"role": message.role, "content": content}
        if message.tool_calls:
            msg["tool_calls"] = message.tool_calls
        if message.tool_results:
            msg["tool_results"] = message.tool_results
        return msg

    @staticmethod
    def _is_empty_assistant_message(message) -> bool:
        """Return True if *message* is an assistant message with no content and no tool_calls.

        Reasoning is not counted as content because most providers (including
        OpenAI) reject assistant messages that only contain reasoning.
        """
        if message.role != "assistant":
            return False
        return (
            not (message.content or "").strip()
            and not message.tool_calls
        )

    def _build_messages(self, user_message: str | None = None) -> list[dict[str, Any]]:
        messages: list[dict[str, Any]] = [
            {"role": "system", "content": self.system_prompt},
        ]
        for msg in self._chat.messages.order_by("created_at"):
            if self._is_empty_assistant_message(msg):
                continue
            messages.append(self._db_message_to_llm(msg))
        if user_message:
            messages.append({"role": "user", "content": user_message})
        provider = self._chat.provider or _infer_provider(self._chat.model)
        messages = _strip_unsupported_file_blocks(messages, provider)
        return _strip_files_from_history(messages)

    def _save_assistant_message(
        self,
        content: str = "",
        tool_calls: list[dict[str, Any]] | None = None,
        reasoning: str = "",
    ) -> Any:
        """Create and return an assistant Message row."""
        return self._chat.messages.create(
            role="assistant",
            content=content,
            tool_calls=tool_calls or None,
            reasoning=reasoning,
        )

    def chat(
        self,
        user_message: str,
        *,
        stream: bool = False,
        reasoning_effort: ReasoningEffort | None = None,
    ) -> str | Iterator[StreamEvent]:
        """Override to persist the user message first."""
        self._chat.messages.create(role="user", content=user_message)
        self._chat.save(update_fields=["updated_at"])
        if stream:
            return self._chat_stream_persisted(reasoning_effort=reasoning_effort)
        return self._chat_sync_persisted(reasoning_effort=reasoning_effort)

    def _chat_sync_persisted(
        self, *, reasoning_effort: ReasoningEffort | None = None
    ) -> str:
        """Non-streaming chat that saves the assistant response to the DB."""
        messages = self._build_messages()

        for _round in range(self.max_tool_rounds):
            response: ChatCompletion = completion(
                model=_resolve_model(self.model, self.provider),
                messages=messages,
                tools=self._tools,
                reasoning_effort=reasoning_effort,
                api_key=self.api_key,
                api_base=self.api_base,
                temperature=self.temperature,
                **self.extra_kwargs,
            )

            msg = self._extract_assistant_message(response)
            self._save_assistant_message(
                content=msg.content or "",
                tool_calls=[
                    {
                        "id": tc.id,
                        "type": tc.type,
                        "function": {
                            "name": tc.function.name,
                            "arguments": tc.function.arguments,
                        },
                    }
                    for tc in (msg.tool_calls or [])
                ],
                reasoning=msg.reasoning.content if msg.reasoning else "",
            )

            llm_msg = self._message_to_dict(msg)
            messages.append(llm_msg)

            if not msg.tool_calls:
                return msg.content or ""

            tool_results = self._run_tool_calls(llm_msg.get("tool_calls", []))
            for tr in tool_results:
                self._chat.messages.create(
                    role="tool",
                    content=tr["content"],
                    tool_results={
                        "tool_call_id": tr["tool_call_id"],
                        "name": tr["name"],
                    },
                )
            messages.extend(tool_results)

        return msg.content or ""

    def _chat_stream_persisted(
        self, *, reasoning_effort: ReasoningEffort | None = None
    ) -> Iterator[StreamEvent]:
        """Streaming chat that persists the assistant response incrementally.

        Yields the same events as the base agent, but also saves content to
        the database as it arrives.
        """
        messages = self._build_messages()
        assistant_db = self._save_assistant_message()

        for _round in range(self.max_tool_rounds):
            chunks: list[ChatCompletionChunk] = []
            stream = completion(
                model=_resolve_model(self.model, self.provider),
                messages=messages,
                tools=self._tools,
                stream=True,
                reasoning_effort=reasoning_effort,
                api_key=self.api_key,
                api_base=self.api_base,
                temperature=self.temperature,
                **self.extra_kwargs,
            )

            for chunk in stream:
                chunks.append(chunk)
                choice = chunk.choices[0] if chunk.choices else None
                if not choice:
                    continue

                delta = choice.delta
                if delta.content:
                    assistant_db.content += delta.content
                    assistant_db.save(update_fields=["content"])
                    yield StreamEvent(type=StreamEventType.CONTENT, text=delta.content)
                if delta.reasoning and delta.reasoning.content:
                    assistant_db.reasoning += delta.reasoning.content
                    assistant_db.save(update_fields=["reasoning"])
                    yield StreamEvent(
                        type=StreamEventType.REASONING,
                        text=delta.reasoning.content,
                    )

            tool_calls = _collect_streaming_tool_calls(chunks)
            content_parts = [
                chunk.choices[0].delta.content
                for chunk in chunks
                if chunk.choices
                and chunk.choices[0].delta
                and chunk.choices[0].delta.content
            ]
            content = "".join(content_parts)

            assistant_db.tool_calls = tool_calls or None
            assistant_db.save(update_fields=["tool_calls"])

            llm_msg: dict[str, Any] = {"role": "assistant", "content": content}
            if tool_calls:
                llm_msg["tool_calls"] = tool_calls
            messages.append(llm_msg)

            if not tool_calls:
                break

            for tc in tool_calls:
                name = tc["function"]["name"]
                try:
                    args = json.loads(tc["function"]["arguments"])
                except json.JSONDecodeError:
                    args = {}
                yield StreamEvent(
                    type=StreamEventType.TOOL_CALL,
                    tool_name=name,
                    tool_args=args,
                )

            tool_results = self._run_tool_calls(tool_calls)
            for tc, result_msg in zip(tool_calls, tool_results):
                yield StreamEvent(
                    type=StreamEventType.TOOL_RESULT,
                    tool_name=tc["function"]["name"],
                    tool_result=result_msg["content"],
                )
                self._chat.messages.create(
                    role="tool",
                    content=result_msg["content"],
                    tool_results={
                        "tool_call_id": result_msg["tool_call_id"],
                        "name": result_msg["name"],
                    },
                )

            messages.extend(tool_results)
            assistant_db = self._save_assistant_message()

        yield StreamEvent(type=StreamEventType.DONE)
