"""OpenAI Completions API provider implementation."""

from __future__ import annotations

import asyncio
import time
from typing import Any

from pi_agents.ai.env_api_keys import get_env_api_key
from pi_agents.ai.models import calculate_cost
from pi_agents.ai.types import (
    AssistantMessage,
    Context,
    Cost,
    Model,
    SimpleStreamOptions,
    TextContent,
    ThinkingContent,
    ToolCall,
    Usage,
)
from pi_agents.ai.utils.event_stream import AssistantMessageEventStream
from pi_agents.ai.utils.json_parse import parse_streaming_json


def _get_api_key_for_provider(model: Model, options: dict[str, Any]) -> str | None:
    """Get API key based on provider."""
    key = options.get("api_key") or options.get("apiKey")
    if key:
        return key

    provider = model.provider
    return get_env_api_key(provider)


def _convert_messages(context: Context) -> list[dict[str, Any]]:
    """Convert pi-agents messages to OpenAI format."""
    messages: list[dict[str, Any]] = []

    if context.system_prompt:
        messages.append({"role": "system", "content": context.system_prompt})

    for msg in context.messages:
        if hasattr(msg, "role"):
            role = msg.role if isinstance(msg.role, str) else msg.get("role", "user")
        elif isinstance(msg, dict):
            role = msg.get("role", "user")
        else:
            continue

        if role == "user":
            content = msg.content if hasattr(msg, "content") else msg.get("content", "")
            if isinstance(content, str):
                messages.append({"role": "user", "content": content})
            else:
                parts = []
                for block in content:
                    btype = block.type if hasattr(block, "type") else block.get("type", "text")
                    if btype == "text":
                        text = block.text if hasattr(block, "text") else block.get("text", "")
                        parts.append({"type": "text", "text": text})
                    elif btype == "image":
                        data = block.data if hasattr(block, "data") else block.get("data", "")
                        mime = block.mime_type if hasattr(block, "mime_type") else block.get("mimeType", "")
                        parts.append({
                            "type": "image_url",
                            "image_url": {"url": f"data:{mime};base64,{data}"},
                        })
                messages.append({"role": "user", "content": parts})

        elif role == "assistant":
            content_list = msg.content if hasattr(msg, "content") else msg.get("content", [])
            text_parts = []
            tool_calls = []
            for block in content_list:
                btype = block.type if hasattr(block, "type") else block.get("type", "text")
                if btype == "text":
                    text = block.text if hasattr(block, "text") else block.get("text", "")
                    text_parts.append(text)
                elif btype == "toolCall":
                    tc_id = block.id if hasattr(block, "id") else block.get("id", "")
                    tc_name = block.name if hasattr(block, "name") else block.get("name", "")
                    tc_args = block.arguments if hasattr(block, "arguments") else block.get("arguments", {})
                    import json
                    tool_calls.append({
                        "id": tc_id,
                        "type": "function",
                        "function": {"name": tc_name, "arguments": json.dumps(tc_args)},
                    })

            entry: dict[str, Any] = {"role": "assistant"}
            if text_parts:
                entry["content"] = "\n".join(text_parts)
            if tool_calls:
                entry["tool_calls"] = tool_calls
            if not text_parts and not tool_calls:
                entry["content"] = ""
            messages.append(entry)

        elif role == "toolResult":
            tc_id = msg.tool_call_id if hasattr(msg, "tool_call_id") else msg.get("toolCallId", "")
            content_list = msg.content if hasattr(msg, "content") else msg.get("content", [])
            text = ""
            for block in content_list:
                btype = block.type if hasattr(block, "type") else block.get("type", "text")
                if btype == "text":
                    text += block.text if hasattr(block, "text") else block.get("text", "")
            messages.append({"role": "tool", "tool_call_id": tc_id, "content": text})

    return messages


def _convert_tools(tools: list[Any] | None) -> list[dict[str, Any]] | None:
    """Convert tools to OpenAI format."""
    if not tools:
        return None
    result = []
    for t in tools:
        name = t.name if hasattr(t, "name") else t.get("name", "")
        desc = t.description if hasattr(t, "description") else t.get("description", "")
        params = t.parameters if hasattr(t, "parameters") else t.get("parameters", {})
        result.append({
            "type": "function",
            "function": {"name": name, "description": desc, "parameters": params},
        })
    return result


def stream_openai_completions(
    model: Model,
    context: Context,
    options: dict[str, Any] | None = None,
) -> AssistantMessageEventStream:
    """Stream responses from the OpenAI Completions API."""
    event_stream = AssistantMessageEventStream()
    opts = options or {}

    async def _run():
        try:
            import openai

            api_key = _get_api_key_for_provider(model, opts)
            if not api_key:
                raise ValueError(f"API key not found for provider {model.provider}")

            client = openai.AsyncOpenAI(
                api_key=api_key,
                base_url=model.base_url or None,
            )

            messages = _convert_messages(context)
            tools = _convert_tools(context.tools)

            params: dict[str, Any] = {
                "model": model.id,
                "messages": messages,
                "stream": True,
                "stream_options": {"include_usage": True},
            }

            max_tokens = opts.get("max_tokens") or opts.get("maxTokens") or model.max_tokens
            if max_tokens:
                params["max_completion_tokens"] = max_tokens

            if tools:
                params["tools"] = tools

            if opts.get("temperature") is not None:
                params["temperature"] = opts["temperature"]

            partial = AssistantMessage(
                role="assistant",
                content=[],
                api=model.api,
                provider=model.provider,
                model=model.id,
                usage=Usage(cost=Cost()),
                stop_reason="stop",
                timestamp=int(time.time() * 1000),
            )

            event_stream.push({"type": "start", "partial": partial})

            text_buffer = ""
            tool_call_buffers: dict[int, dict[str, Any]] = {}
            content_index = 0

            stream = await client.chat.completions.create(**params)
            async for chunk in stream:
                if not chunk.choices:
                    # Usage chunk
                    if chunk.usage:
                        partial.usage.input = chunk.usage.prompt_tokens or 0
                        partial.usage.output = chunk.usage.completion_tokens or 0
                        partial.usage.total_tokens = chunk.usage.total_tokens or 0
                    continue

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

                # Text content
                if delta and delta.content:
                    if not text_buffer:
                        event_stream.push({"type": "text_start", "contentIndex": content_index, "partial": partial})
                    text_buffer += delta.content
                    event_stream.push({
                        "type": "text_delta",
                        "contentIndex": content_index,
                        "delta": delta.content,
                        "partial": partial,
                    })

                # Tool calls
                if delta and delta.tool_calls:
                    for tc in delta.tool_calls:
                        idx = tc.index
                        if idx not in tool_call_buffers:
                            if text_buffer:
                                event_stream.push({
                                    "type": "text_end",
                                    "contentIndex": content_index,
                                    "content": text_buffer,
                                    "partial": partial,
                                })
                                content_index += 1
                                text_buffer = ""
                            tool_call_buffers[idx] = {
                                "id": tc.id or "",
                                "name": tc.function.name if tc.function and tc.function.name else "",
                                "arguments": "",
                            }
                            event_stream.push({
                                "type": "toolcall_start",
                                "contentIndex": content_index + idx,
                                "partial": partial,
                            })
                        if tc.function and tc.function.arguments:
                            tool_call_buffers[idx]["arguments"] += tc.function.arguments
                            event_stream.push({
                                "type": "toolcall_delta",
                                "contentIndex": content_index + idx,
                                "delta": tc.function.arguments,
                                "partial": partial,
                            })

                # Finish
                if choice.finish_reason:
                    if text_buffer:
                        event_stream.push({
                            "type": "text_end",
                            "contentIndex": content_index,
                            "content": text_buffer,
                            "partial": partial,
                        })
                        partial.content.append(TextContent(type="text", text=text_buffer))

                    for idx, tc_buf in sorted(tool_call_buffers.items()):
                        args = parse_streaming_json(tc_buf["arguments"])
                        tc = ToolCall(
                            type="toolCall",
                            id=tc_buf["id"],
                            name=tc_buf["name"],
                            arguments=args,
                        )
                        partial.content.append(tc)
                        event_stream.push({
                            "type": "toolcall_end",
                            "contentIndex": content_index + idx,
                            "toolCall": tc,
                            "partial": partial,
                        })

                    reason_map = {"stop": "stop", "length": "length", "tool_calls": "toolUse"}
                    partial.stop_reason = reason_map.get(choice.finish_reason, "stop")

            calculate_cost(model, partial.usage)
            event_stream.push({"type": "done", "reason": partial.stop_reason, "message": partial})

        except Exception as e:
            error_msg = AssistantMessage(
                role="assistant",
                content=[],
                api=model.api,
                provider=model.provider,
                model=model.id,
                usage=Usage(cost=Cost()),
                stop_reason="error",
                error_message=str(e),
                timestamp=int(time.time() * 1000),
            )
            event_stream.push({"type": "error", "reason": "error", "error": error_msg})

        event_stream.end()

    asyncio.ensure_future(_run())
    return event_stream


def stream_simple_openai_completions(
    model: Model,
    context: Context,
    options: SimpleStreamOptions | None = None,
) -> AssistantMessageEventStream:
    """Stream responses using simplified options."""
    opts: dict[str, Any] = {}
    if options:
        if options.temperature is not None:
            opts["temperature"] = options.temperature
        if options.max_tokens is not None:
            opts["maxTokens"] = options.max_tokens
        if options.api_key:
            opts["apiKey"] = options.api_key
    return stream_openai_completions(model, context, opts)
