"""OpenAI provider for Logos ADK."""

from __future__ import annotations

import base64
import mimetypes
import os
from pathlib import Path
import httpx
from typing import Any, AsyncIterator

from logos_adk.errors import AuthenticationError, ProviderError
from logos_adk.providers import Provider, provider_registry
from logos_adk.types import (
    ContentPart,
    Message,
    MediaChunk,
    Response,
    TextChunk,
    ToolCall,
    Usage,
    content_text,
    normalize_content_parts,
)


def _image_url_for_part(part: ContentPart) -> str | None:
    if part.uri:
        if part.uri.startswith(("http://", "https://", "data:")):
            return part.uri
        candidate = Path(part.uri)
        if candidate.exists() and candidate.is_file():
            mime_type = part.mime_type or mimetypes.guess_type(candidate.name)[0] or "application/octet-stream"
            return f"data:{mime_type};base64,{base64.b64encode(candidate.read_bytes()).decode('ascii')}"
    if part.data and part.mime_type:
        return f"data:{part.mime_type};base64,{part.data}"
    return None


def _content_to_openai(message: Message) -> str | list[dict[str, Any]] | None:
    parts = normalize_content_parts(message.content_parts)
    if not parts:
        return message.content

    api_parts: list[dict[str, Any]] = []
    if message.content and not any(part.type == "text" and part.text == message.content for part in parts):
        api_parts.append({"type": "text", "text": message.content})

    for part in parts:
        if part.type == "text":
            api_parts.append({"type": "text", "text": part.text or ""})
            continue
        if part.type == "image":
            image_url = _image_url_for_part(part)
            if image_url is not None:
                payload: dict[str, Any] = {"url": image_url}
                detail = part.metadata.get("detail") if isinstance(part.metadata, dict) else None
                if isinstance(detail, str) and detail:
                    payload["detail"] = detail
                api_parts.append({"type": "image_url", "image_url": payload})
                continue
        fallback_label = part.name or part.mime_type or part.uri or part.type
        api_parts.append({"type": "text", "text": f"[{part.type}: {fallback_label}]"})

    return api_parts or None


def _messages_to_openai(messages: list[Message]) -> list[dict]:
    """Convert Logos messages to OpenAI API format."""
    api_messages: list[dict[str, Any]] = []

    for msg in messages:
        if msg.role == "tool":
            if msg.tool_result is not None:
                api_messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": msg.tool_result.call_id,
                        "content": msg.tool_result.error
                        or content_text(
                            str(msg.tool_result.content) if msg.tool_result.content is not None else "",
                            msg.tool_result.content_parts,
                        ),
                    }
                )
            continue

        if msg.role == "assistant" and msg.tool_calls:
            tool_calls_api = []
            for tc in msg.tool_calls:
                import json

                tool_calls_api.append(
                    {
                        "id": tc.id,
                        "type": "function",
                        "function": {
                            "name": tc.name,
                            "arguments": json.dumps(tc.arguments),
                        },
                    }
                )
            entry: dict[str, Any] = {
                "role": "assistant",
                "content": _content_to_openai(msg),
                "tool_calls": tool_calls_api,
            }
            api_messages.append(entry)
            continue

        api_messages.append({"role": msg.role, "content": _content_to_openai(msg)})

    return api_messages


def _tools_to_openai(tools: list[dict] | None) -> list[dict] | None:
    """Convert Logos tool schemas to OpenAI function-calling format."""
    if not tools:
        return None
    return [
        {
            "type": "function",
            "function": {
                "name": t["name"],
                "description": t.get("description", ""),
                "parameters": t.get("parameters", {"type": "object", "properties": {}}),
            },
        }
        for t in tools
    ]


_INTERNAL_CREATE_KWARG_KEYS = {
    "campaign_id",
    "correlation_id",
    "experiment_id",
    "max_events",
    "on_done",
    "on_text",
    "on_tool_end",
    "on_tool_start",
    "payload",
    "prompt_name",
    "prompt_variant",
    "prompt_version",
    "publisher",
    "request_id",
    "routing_hint",
    "run_id",
    "session_id",
    "task_type",
    "user_segment",
    "workspace_id",
}


def _filter_openai_create_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
    """Drop internal ADK/runtime kwargs before calling the OpenAI SDK."""
    return {
        key: value
        for key, value in kwargs.items()
        if not key.startswith("_") and key not in _INTERNAL_CREATE_KWARG_KEYS
    }


def _parse_message_content(raw_content: Any) -> tuple[str, list[ContentPart]]:
    if isinstance(raw_content, str):
        return raw_content, []
    if not isinstance(raw_content, list):
        return "", []

    parts: list[ContentPart] = []
    for item in raw_content:
        if not isinstance(item, dict):
            continue
        item_type = str(item.get("type") or "text")
        if item_type == "text":
            parts.append(ContentPart(type="text", text=item.get("text") or ""))
            continue
        if item_type == "image_url":
            payload = item.get("image_url") or {}
            if isinstance(payload, dict):
                metadata = {}
                if payload.get("detail"):
                    metadata["detail"] = payload.get("detail")
                parts.append(ContentPart(type="image", uri=payload.get("url"), metadata=metadata))
            continue
        parts.append(ContentPart(type=item_type, metadata=dict(item)))
    return content_text("", parts), parts


def _parse_response(raw: Any) -> Response:
    """Parse an OpenAI response into a Logos Response."""
    choice = raw.choices[0]
    message = choice.message

    tool_calls: list[ToolCall] = []
    if message.tool_calls:
        import json

        for tc in message.tool_calls:
            tool_calls.append(
                ToolCall(
                    id=tc.id,
                    name=tc.function.name,
                    arguments=json.loads(tc.function.arguments),
                )
            )

    reasoning = None
    if hasattr(message, "reasoning_content") and message.reasoning_content:
        reasoning = message.reasoning_content
    elif hasattr(raw, "extra_data") and raw.extra_data:
        reasoning = raw.extra_data.get("reasoning")
    elif hasattr(message, "model_extra"):
        reasoning = message.model_extra.get("reasoning") if message.model_extra else None

    parsed_content, parsed_parts = _parse_message_content(message.content)

    return Response(
        content=message.content if isinstance(message.content, str) else parsed_content,
        content_parts=parsed_parts,
        tool_calls=tool_calls,
        usage=Usage(
            input_tokens=raw.usage.prompt_tokens,
            output_tokens=raw.usage.completion_tokens,
        ),
        raw={"id": raw.id, "model": raw.model, "finish_reason": choice.finish_reason},
        reasoning=reasoning,
    )


class OpenAIProvider(Provider):
    """Provider for OpenAI models (GPT-4, GPT-4o, etc.)."""

    _http_client_pool: dict[str | None, httpx.AsyncClient] = {}
    _client_pool: dict[tuple[str, str | None, str | None], Any] = {}

    @classmethod
    def _shared_http_client(cls, proxy: str | None) -> httpx.AsyncClient:
        if proxy not in cls._http_client_pool:
            kwargs: dict[str, Any] = {"timeout": 120.0}
            if proxy:
                kwargs["proxy"] = proxy
            cls._http_client_pool[proxy] = httpx.AsyncClient(**kwargs)
        return cls._http_client_pool[proxy]

    def __init__(
        self,
        model: str = "gpt-4o",
        *,
        api_key: str | None = None,
        base_url: str | None = None,
        proxy: str | None = None,
    ) -> None:
        try:
            import openai
        except ImportError:
            raise ImportError(
                "openai package required. Install with: pip install logos-adk[openai]"
            )

        resolved_key = api_key or os.environ.get("OPENAI_API_KEY")
        if not resolved_key:
            raise AuthenticationError(
                "OpenAI API key not found. Set OPENAI_API_KEY env var "
                "or pass api_key= to the provider."
            )

        client_key = (resolved_key, base_url, proxy)
        if client_key not in self._client_pool:
            self._client_pool[client_key] = openai.AsyncOpenAI(
                api_key=resolved_key,
                base_url=base_url,
                http_client=self._shared_http_client(proxy),
                max_retries=2,
            )
        self._client = self._client_pool[client_key]
        self._model = model
        self.proxy = proxy

    async def generate(
        self,
        messages: list[Message],
        tools: list[dict] | None = None,
        output_schema: Any = None,
        **kwargs: Any,
    ) -> Response:
        api_messages = _messages_to_openai(messages)
        api_tools = _tools_to_openai(tools)

        create_kwargs: dict[str, Any] = {
            "model": self._model,
            "messages": api_messages,
        }
        if api_tools:
            create_kwargs["tools"] = api_tools
        create_kwargs.update(_filter_openai_create_kwargs(kwargs))

        try:
            raw = await self._client.chat.completions.create(**create_kwargs)
        except Exception as e:
            raise ProviderError(f"OpenAI API error: {e}") from e

        return _parse_response(raw)

    async def stream(
        self,
        messages: list[Message],
        tools: list[dict] | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[Any]:
        api_messages = _messages_to_openai(messages)
        api_tools = _tools_to_openai(tools)

        create_kwargs: dict[str, Any] = {
            "model": self._model,
            "messages": api_messages,
            "stream": True,
        }
        if api_tools:
            create_kwargs["tools"] = api_tools
        create_kwargs.update(_filter_openai_create_kwargs(kwargs))

        try:
            stream = await self._client.chat.completions.create(**create_kwargs)
            collected_text: list[str] = []
            usage = None

            async for chunk in stream:
                if not chunk.choices:
                    continue
                delta = chunk.choices[0].delta
                if isinstance(delta.content, str) and delta.content:
                    collected_text.append(delta.content)
                    yield TextChunk(text=delta.content)
                elif isinstance(delta.content, list):
                    parsed_text, parsed_parts = _parse_message_content(delta.content)
                    if parsed_text:
                        collected_text.append(parsed_text)
                        yield TextChunk(text=parsed_text)
                    for part in parsed_parts:
                        if part.type != "text":
                            yield MediaChunk(part=part)

                # Try to get usage from chunk if available
                if chunk.usage and not usage:
                    usage = chunk.usage

            # Use collected usage or estimate from text length
            if usage:
                input_tokens = usage.prompt_tokens or 0
                output_tokens = usage.completion_tokens or 0
            else:
                # Estimate tokens from text length (rough approximation: ~4 chars per token)
                input_tokens = sum(len(str(m.get("content", ""))) for m in api_messages) // 4
                output_tokens = len("".join(collected_text)) // 4

            yield Response(
                content="".join(collected_text),
                tool_calls=[],
                usage=Usage(input_tokens=input_tokens, output_tokens=output_tokens),
                raw={},
            )
        except Exception as e:
            raise ProviderError(f"OpenAI streaming error: {e}") from e


class OpenAICompatibleProvider(OpenAIProvider):
    """Provider for any OpenAI-compatible API (Ollama, LM Studio, vLLM, Together, etc.).

    Usage:
        # Ollama
        provider = OpenAICompatibleProvider(
            model="llama3",
            base_url="http://localhost:11434/v1",
        )

        # LM Studio
        provider = OpenAICompatibleProvider(
            model="local-model",
            base_url="http://localhost:1234/v1",
        )

        # Together AI
        provider = OpenAICompatibleProvider(
            model="meta-llama/Llama-3-70b-chat-hf",
            base_url="https://api.together.xyz/v1",
            api_key="...",
        )

        # Polza AI
        provider = OpenAICompatibleProvider(
            model="minimax/minimax-m2.5",
            base_url="https://api.polza.ai/api/v1",
            api_key="pza_...",
        )

        agent = Agent("bot", model=provider)
    """

    def __init__(
        self,
        model: str,
        *,
        base_url: str,
        api_key: str | None = None,
        api_key_env: str | None = None,
        proxy: str | None = None,
    ) -> None:
        try:
            import openai
        except ImportError:
            raise ImportError(
                "openai package required. Install with: pip install logos-adk[openai]"
            )

        resolved_key = (
            api_key
            or (os.environ.get(api_key_env) if api_key_env else None)
            or os.environ.get("OPENAI_API_KEY")
        )

        self._auth_header: str | None = None
        if resolved_key:
            self._auth_header = f"Bearer {resolved_key}"

        api_key_for_client = "not-needed" if self._auth_header else (resolved_key or "not-needed")

        client_key = (api_key_for_client, base_url, proxy)
        if client_key not in self._client_pool:
            self._client_pool[client_key] = openai.AsyncOpenAI(
                api_key=api_key_for_client,
                base_url=base_url,
                http_client=self._shared_http_client(proxy),
            )
        self._client = self._client_pool[client_key]
        self._model = model
        self.proxy = proxy

    async def generate(
        self,
        messages: list[Message],
        tools: list[dict] | None = None,
        output_schema: Any = None,
        **kwargs: Any,
    ) -> Response:
        if self._auth_header:
            kwargs["extra_headers"] = {"Authorization": self._auth_header}
        return await super().generate(messages, tools, output_schema, **kwargs)

    async def stream(
        self,
        messages: list[Message],
        tools: list[dict] | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[Any]:
        if self._auth_header:
            kwargs["extra_headers"] = {"Authorization": self._auth_header}
        async for chunk in super().stream(messages, tools, **kwargs):
            yield chunk


def _openai_factory(model_name: str) -> OpenAIProvider:
    return OpenAIProvider(model=model_name)


# Auto-register for gpt-* and o1-* model patterns
provider_registry.register("gpt-*", _openai_factory)
provider_registry.register("o1-*", _openai_factory)
provider_registry.register("o3-*", _openai_factory)
