"""OpenAI Responses 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,
    ToolCall,
    Usage,
)
from pi_agents.ai.utils.event_stream import AssistantMessageEventStream
from pi_agents.ai.utils.json_parse import parse_streaming_json


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

    async def _run():
        try:
            import openai

            api_key = opts.get("api_key") or opts.get("apiKey") or get_env_api_key(model.provider)
            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,
            )

            # Build input from context
            input_messages: list[dict[str, Any]] = []
            if context.system_prompt:
                input_messages.append({"role": "developer", "content": context.system_prompt})

            for msg in context.messages:
                role = msg.role if hasattr(msg, "role") else msg.get("role", "user")
                if role == "user":
                    content = msg.content if hasattr(msg, "content") else msg.get("content", "")
                    input_messages.append({"role": "user", "content": content if isinstance(content, str) else str(content)})
                elif role == "assistant":
                    content_list = msg.content if hasattr(msg, "content") else msg.get("content", [])
                    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", "")
                            input_messages.append({"role": "assistant", "content": text})

            params: dict[str, Any] = {
                "model": model.id,
                "input": input_messages,
                "stream": True,
            }

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

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

            # Tools
            if context.tools:
                params["tools"] = [
                    {
                        "type": "function",
                        "name": t.name if hasattr(t, "name") else t.get("name", ""),
                        "description": t.description if hasattr(t, "description") else t.get("description", ""),
                        "parameters": t.parameters if hasattr(t, "parameters") else t.get("parameters", {}),
                    }
                    for t in context.tools
                ]

            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})

            # Use responses API
            stream = await client.responses.create(**params)
            text_buffer = ""
            content_index = 0

            async for event in stream:
                event_type = event.type if hasattr(event, "type") else ""

                if event_type == "response.output_text.delta":
                    delta = event.delta if hasattr(event, "delta") else ""
                    if not text_buffer:
                        event_stream.push({"type": "text_start", "contentIndex": content_index, "partial": partial})
                    text_buffer += delta
                    event_stream.push({
                        "type": "text_delta",
                        "contentIndex": content_index,
                        "delta": delta,
                        "partial": partial,
                    })

                elif event_type == "response.output_text.done":
                    text = event.text if hasattr(event, "text") else text_buffer
                    event_stream.push({
                        "type": "text_end",
                        "contentIndex": content_index,
                        "content": text,
                        "partial": partial,
                    })
                    partial.content.append(TextContent(type="text", text=text))
                    text_buffer = ""
                    content_index += 1

                elif event_type == "response.completed":
                    response = event.response if hasattr(event, "response") else None
                    if response and hasattr(response, "usage"):
                        u = response.usage
                        partial.usage.input = getattr(u, "input_tokens", 0)
                        partial.usage.output = getattr(u, "output_tokens", 0)
                        partial.usage.total_tokens = partial.usage.input + partial.usage.output

            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_responses(
    model: Model,
    context: Context,
    options: SimpleStreamOptions | None = None,
) -> AssistantMessageEventStream:
    """Stream with 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_responses(model, context, opts)
