"""
providers/openrouter_provider.py — OpenRouter provider.

OpenRouter exposes an OpenAI-compatible API, so we use the `openai` Python SDK
with a custom base_url. The X-Title header identifies the app to OpenRouter
(shows as "oAI-Web" in OpenRouter usage logs).

Tool schemas need conversion: oAI-Web uses Anthropic-native format internally,
OpenRouter expects OpenAI function-calling format.
"""
from __future__ import annotations

import json
import logging
from typing import Any

from openai import OpenAI, AsyncOpenAI

from .base import AIProvider, ProviderResponse, ToolCallResult, UsageStats

logger = logging.getLogger(__name__)

OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"
DEFAULT_MODEL = "anthropic/claude-sonnet-4-5"


def _convert_content_blocks(blocks: list[dict]) -> list[dict]:
    """Convert Anthropic-native content blocks to OpenAI image_url / file format."""
    result = []
    for block in blocks:
        btype = block.get("type")
        if btype in ("image", "document"):
            src = block.get("source", {})
            if src.get("type") == "base64":
                data_url = f"data:{src['media_type']};base64,{src['data']}"
                result.append({"type": "image_url", "image_url": {"url": data_url}})
        else:
            result.append(block)
    return result


class OpenRouterProvider(AIProvider):
    def __init__(self, api_key: str, app_name: str = "oAI-Web", app_url: str = "https://mac.oai.pm") -> None:
        extra_headers = {
            "X-Title": app_name,
            "HTTP-Referer": app_url,
        }

        self._client = OpenAI(
            api_key=api_key,
            base_url=OPENROUTER_BASE_URL,
            default_headers=extra_headers,
        )
        self._async_client = AsyncOpenAI(
            api_key=api_key,
            base_url=OPENROUTER_BASE_URL,
            default_headers=extra_headers,
        )

    @property
    def name(self) -> str:
        return "OpenRouter"

    @property
    def default_model(self) -> str:
        return DEFAULT_MODEL

    # ── Public interface ──────────────────────────────────────────────────────

    def chat(
        self,
        messages: list[dict],
        tools: list[dict] | None = None,
        system: str = "",
        model: str = "",
        max_tokens: int = 4096,
    ) -> ProviderResponse:
        params = self._build_params(messages, tools, system, model, max_tokens)
        try:
            response = self._client.chat.completions.create(**params)
            return self._parse_response(response)
        except Exception as e:
            logger.error(f"OpenRouter chat error: {e}")
            return ProviderResponse(text=f"Error: {e}", finish_reason="error")

    async def chat_async(
        self,
        messages: list[dict],
        tools: list[dict] | None = None,
        system: str = "",
        model: str = "",
        max_tokens: int = 4096,
    ) -> ProviderResponse:
        params = self._build_params(messages, tools, system, model, max_tokens)
        try:
            response = await self._async_client.chat.completions.create(**params)
            return self._parse_response(response)
        except Exception as e:
            logger.error(f"OpenRouter async chat error: {e}")
            return ProviderResponse(text=f"Error: {e}", finish_reason="error")

    # ── Internal helpers ──────────────────────────────────────────────────────

    def _build_params(
        self,
        messages: list[dict],
        tools: list[dict] | None,
        system: str,
        model: str,
        max_tokens: int,
    ) -> dict:
        effective_model = model or self.default_model

        # Detect image-generation models via output_modalities in the OR cache
        from .models import get_or_output_modalities
        bare_id = effective_model.removeprefix("openrouter:")
        out_modalities = get_or_output_modalities(bare_id)
        is_image_gen = "image" in out_modalities

        openai_messages = self._convert_messages(messages, system)
        params: dict = {"model": effective_model, "messages": openai_messages}

        if is_image_gen:
            # Image-gen models use modalities parameter; max_tokens not applicable
            params["modalities"] = out_modalities
        else:
            params["max_tokens"] = max_tokens
            if tools:
                params["tools"] = [self._to_openai_tool(t) for t in tools]
                params["tool_choice"] = "auto"
        return params

    def _convert_messages(self, messages: list[dict], system: str) -> list[dict]:
        """
        Convert aide's internal message list to OpenAI format.
        Prepend system message if provided.

        aide internal format uses:
          - assistant with "tool_calls": [{"id", "name", "arguments"}]
          - role "tool" with "tool_call_id" and "content"

        OpenAI format uses:
          - assistant with "tool_calls": [{"id", "type": "function", "function": {"name", "arguments"}}]
          - role "tool" with "tool_call_id" and "content"
        """
        result: list[dict] = []

        if system:
            result.append({"role": "system", "content": system})

        i = 0
        while i < len(messages):
            msg = messages[i]
            role = msg["role"]

            if role == "system":
                i += 1
                continue  # Already prepended above

            if role == "assistant" and msg.get("tool_calls"):
                openai_tool_calls = []
                for tc in msg["tool_calls"]:
                    openai_tool_calls.append({
                        "id": tc["id"],
                        "type": "function",
                        "function": {
                            "name": tc["name"],
                            "arguments": json.dumps(tc["arguments"]),
                        },
                    })
                out: dict[str, Any] = {"role": "assistant", "tool_calls": openai_tool_calls}
                if msg.get("content"):
                    out["content"] = msg["content"]
                result.append(out)

            elif role == "tool":
                # Group consecutive tool results; collect any image blocks for injection
                pending_images: list[dict] = []
                while i < len(messages) and messages[i]["role"] == "tool":
                    t = messages[i]
                    content = t.get("content", "")
                    if isinstance(content, list):
                        text = " ".join(b.get("text", "") for b in content if b.get("type") == "text") or "[image]"
                        pending_images.extend(b for b in content if b.get("type") == "image")
                        content = text
                    result.append({"role": "tool", "tool_call_id": t["tool_call_id"], "content": content})
                    i += 1
                if pending_images:
                    result.append({"role": "user", "content": _convert_content_blocks(pending_images)})
                continue  # i already advanced

            else:
                content = msg.get("content", "")
                if isinstance(content, list):
                    content = _convert_content_blocks(content)
                result.append({"role": role, "content": content})

            i += 1

        return result

    @staticmethod
    def _to_openai_tool(aide_tool: dict) -> dict:
        """
        Convert aide's tool schema (Anthropic-native) to OpenAI function-calling format.

        Anthropic format:
          {"name": "...", "description": "...", "input_schema": {...}}

        OpenAI format:
          {"type": "function", "function": {"name": "...", "description": "...", "parameters": {...}}}
        """
        return {
            "type": "function",
            "function": {
                "name": aide_tool["name"],
                "description": aide_tool.get("description", ""),
                "parameters": aide_tool.get("input_schema", {"type": "object", "properties": {}}),
            },
        }

    def _parse_response(self, response) -> ProviderResponse:
        choice = response.choices[0] if response.choices else None
        if not choice:
            return ProviderResponse(text=None, finish_reason="error")

        message = choice.message
        text = message.content or None
        tool_calls: list[ToolCallResult] = []

        if message.tool_calls:
            for tc in message.tool_calls:
                try:
                    arguments = json.loads(tc.function.arguments)
                except json.JSONDecodeError:
                    arguments = {"_raw": tc.function.arguments}
                tool_calls.append(ToolCallResult(
                    id=tc.id,
                    name=tc.function.name,
                    arguments=arguments,
                ))

        usage = UsageStats()
        if response.usage:
            usage = UsageStats(
                input_tokens=response.usage.prompt_tokens,
                output_tokens=response.usage.completion_tokens,
            )

        finish_reason = choice.finish_reason or "stop"
        if tool_calls:
            finish_reason = "tool_use"

        # Extract generated images.
        # OpenRouter image structure: {"image_url": {"url": "data:image/png;base64,..."}}
        # Two possible locations (both checked; first non-empty wins):
        #   A. message.images  — top-level field in the message (custom OpenRouter format)
        #   B. message.content — array of content blocks with type "image_url"
        images: list[str] = []

        def _url_from_img_obj(img) -> str:
            """Extract URL string from an image object in OpenRouter format."""
            if isinstance(img, str):
                return img
            if isinstance(img, dict):
                # {"image_url": {"url": "..."}}  ← OpenRouter format
                inner = img.get("image_url")
                if isinstance(inner, dict):
                    return inner.get("url") or ""
                # Fallback: {"url": "..."}
                return img.get("url") or ""
            # Pydantic model object with image_url attribute
            image_url_obj = getattr(img, "image_url", None)
            if image_url_obj is not None:
                return getattr(image_url_obj, "url", None) or ""
            return ""

        # A. message.model_extra["images"] (SDK stores unknown fields here)
        extra = getattr(message, "model_extra", None) or {}
        raw_images = extra.get("images") or getattr(message, "images", None) or []
        for img in raw_images:
            url = _url_from_img_obj(img)
            if url:
                images.append(url)

        # B. Content as array of blocks: [{"type":"image_url","image_url":{"url":"..."}}]
        if not images:
            raw_content = message.content
            if isinstance(raw_content, list):
                for block in raw_content:
                    if isinstance(block, dict) and block.get("type") == "image_url":
                        url = (block.get("image_url") or {}).get("url") or ""
                        if url:
                            images.append(url)

        logger.info("[openrouter] image-gen response: %d image(s), text=%r, extra_keys=%s",
                    len(images), text[:80] if text else None, list(extra.keys()))

        return ProviderResponse(
            text=text,
            tool_calls=tool_calls,
            usage=usage,
            finish_reason=finish_reason,
            model=response.model,
            images=images,
        )
