"""OpenAI-compatible provider (GPT, Ollama, etc.) with streaming support."""

from __future__ import annotations

import json
from typing import Any, Generator

import openai

from tsumugi.providers.base import (
    AssistantMessage,
    ContentBlock,
    ProviderBase,
    StreamEvent,
    StreamEventType,
    TextBlock,
    ToolCallBlock,
    ToolDefinition,
    Usage,
)


class OpenAICompatProvider(ProviderBase):
    """Provider for OpenAI-compatible APIs (GPT, Ollama, etc.)."""

    def __init__(
        self,
        api_key: str,
        model: str = "gpt-4o",
        base_url: str | None = None,
    ):
        kwargs: dict[str, Any] = {"api_key": api_key}
        if base_url:
            kwargs["base_url"] = base_url
        self.client = openai.OpenAI(**kwargs)
        self.model = model
        self.base_url = base_url

    def _convert_tools(
        self, tools: list[ToolDefinition] | None
    ) -> list[dict[str, Any]] | openai.NOT_GIVEN:
        """Convert tool definitions to OpenAI function calling format."""
        if not tools:
            return openai.NOT_GIVEN
        return [
            {
                "type": "function",
                "function": {
                    "name": t.name,
                    "description": t.description,
                    "parameters": t.parameters,
                },
            }
            for t in tools
        ]

    def _convert_messages(
        self, messages: list[dict[str, Any]], system: str
    ) -> list[dict[str, Any]]:
        """Convert Anthropic-format messages to OpenAI format."""
        oai_messages: list[dict[str, Any]] = [
            {"role": "system", "content": system}
        ]

        for msg in messages:
            role = msg["role"]
            content = msg["content"]

            if role == "user":
                if isinstance(content, str):
                    oai_messages.append({"role": "user", "content": content})
                elif isinstance(content, list):
                    # Could be tool_result blocks
                    for block in content:
                        if isinstance(block, dict) and block.get("type") == "tool_result":
                            oai_messages.append({
                                "role": "tool",
                                "tool_call_id": block["tool_use_id"],
                                "content": block.get("content", ""),
                            })
                        elif isinstance(block, dict) and block.get("type") == "text":
                            oai_messages.append({
                                "role": "user",
                                "content": block["text"],
                            })
                        else:
                            oai_messages.append({
                                "role": "user",
                                "content": str(block),
                            })

            elif role == "assistant":
                if isinstance(content, str):
                    oai_messages.append({"role": "assistant", "content": content})
                elif isinstance(content, list):
                    # May contain text and tool_use blocks
                    text_parts = []
                    tool_calls = []
                    for block in content:
                        if isinstance(block, dict):
                            if block.get("type") == "text":
                                text_parts.append(block["text"])
                            elif block.get("type") == "tool_use":
                                tool_calls.append({
                                    "id": block["id"],
                                    "type": "function",
                                    "function": {
                                        "name": block["name"],
                                        "arguments": json.dumps(block["input"]),
                                    },
                                })

                    oai_msg: dict[str, Any] = {"role": "assistant"}
                    if text_parts:
                        oai_msg["content"] = "\n".join(text_parts)
                    else:
                        oai_msg["content"] = None
                    if tool_calls:
                        oai_msg["tool_calls"] = tool_calls
                    oai_messages.append(oai_msg)

        return oai_messages

    def stream(
        self,
        messages: list[dict[str, Any]],
        system: str,
        tools: list[ToolDefinition] | None = None,
        max_tokens: int = 8192,
    ) -> Generator[StreamEvent, None, None]:
        """Stream a response from OpenAI-compatible API."""
        oai_messages = self._convert_messages(messages, system)
        oai_tools = self._convert_tools(tools)

        kwargs: dict[str, Any] = {
            "model": self.model,
            "messages": oai_messages,
            "max_tokens": max_tokens,
            "stream": True,
        }
        if oai_tools is not openai.NOT_GIVEN:
            kwargs["tools"] = oai_tools

        stream = self.client.chat.completions.create(**kwargs)

        # Track tool calls being built up
        current_tool_calls: dict[int, dict[str, str]] = {}
        input_tokens = 0
        output_tokens = 0

        for chunk in stream:
            # Extract usage if available
            if hasattr(chunk, "usage") and chunk.usage:
                input_tokens = getattr(chunk.usage, "prompt_tokens", 0) or 0
                output_tokens = getattr(chunk.usage, "completion_tokens", 0) or 0
            if not chunk.choices:
                continue

            choice = chunk.choices[0]
            delta = choice.delta

            # Text content
            if delta and delta.content:
                yield StreamEvent(
                    type=StreamEventType.TEXT_DELTA,
                    text=delta.content,
                )

            # Tool calls
            if delta and delta.tool_calls:
                for tc in delta.tool_calls:
                    idx = tc.index
                    if idx not in current_tool_calls:
                        current_tool_calls[idx] = {
                            "id": tc.id or "",
                            "name": "",
                            "arguments": "",
                        }
                        if tc.id:
                            current_tool_calls[idx]["id"] = tc.id

                    if tc.function:
                        if tc.function.name:
                            current_tool_calls[idx]["name"] = tc.function.name
                            yield StreamEvent(
                                type=StreamEventType.TOOL_CALL_START,
                                tool_call_id=current_tool_calls[idx]["id"],
                                tool_name=tc.function.name,
                            )
                        if tc.function.arguments:
                            current_tool_calls[idx]["arguments"] += tc.function.arguments
                            yield StreamEvent(
                                type=StreamEventType.TOOL_CALL_DELTA,
                                tool_call_id=current_tool_calls[idx]["id"],
                                tool_name=current_tool_calls[idx]["name"],
                                tool_input_delta=tc.function.arguments,
                            )

            # End of message
            if choice.finish_reason:
                # Emit tool call end events
                for idx, tc_data in current_tool_calls.items():
                    yield StreamEvent(
                        type=StreamEventType.TOOL_CALL_END,
                        tool_call_id=tc_data["id"],
                        tool_name=tc_data["name"],
                        tool_input_delta=tc_data["arguments"],
                    )

                yield StreamEvent(
                    type=StreamEventType.MESSAGE_END,
                    stop_reason=choice.finish_reason,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                )

    def build_assistant_message(
        self, events: list[StreamEvent]
    ) -> AssistantMessage:
        """Reconstruct assistant message from stream events."""
        content: list[ContentBlock] = []
        text_parts: list[str] = []
        tool_inputs: dict[str, str] = {}
        stop_reason = ""
        input_tokens = 0
        output_tokens = 0

        for ev in events:
            if ev.type == StreamEventType.TEXT_DELTA:
                text_parts.append(ev.text)

            elif ev.type == StreamEventType.TOOL_CALL_START:
                if text_parts:
                    content.append(TextBlock(text="".join(text_parts)))
                    text_parts = []
                tool_inputs[ev.tool_call_id] = ""

            elif ev.type == StreamEventType.TOOL_CALL_DELTA:
                tool_inputs[ev.tool_call_id] = (
                    tool_inputs.get(ev.tool_call_id, "") + ev.tool_input_delta
                )

            elif ev.type == StreamEventType.TOOL_CALL_END:
                raw = tool_inputs.get(ev.tool_call_id, ev.tool_input_delta)
                try:
                    parsed = json.loads(raw) if raw else {}
                except json.JSONDecodeError:
                    parsed = {}
                content.append(
                    ToolCallBlock(
                        id=ev.tool_call_id,
                        name=ev.tool_name,
                        input=parsed,
                    )
                )

            elif ev.type == StreamEventType.MESSAGE_END:
                stop_reason = ev.stop_reason
                input_tokens = ev.input_tokens
                output_tokens = ev.output_tokens

        if text_parts:
            content.append(TextBlock(text="".join(text_parts)))

        return AssistantMessage(
            content=content,
            stop_reason=stop_reason,
            usage=Usage(input_tokens=input_tokens, output_tokens=output_tokens),
        )
