"""OpenAI and OpenAI-compatible provider bridge."""

from __future__ import annotations

import json
import logging
from typing import AsyncIterator

import aiohttp

from backend.bridge.base import BaseBridge, BridgeRequest, BridgeResponse, StreamResult, StreamUsage
from backend.config import ModelConfig

logger = logging.getLogger("ecate.bridge.openai")


class OpenAIBridge(BaseBridge):
    """Bridge for OpenAI API and any OpenAI-compatible endpoint (vLLM, Ollama, etc.)."""

    @property
    def provider_type(self) -> str:
        return "openai"

    def _build_url(self) -> str:
        base = self.provider.base_url.rstrip("/")
        if not base.endswith("/v1"):
            base += "/v1" if "/v1" not in base else ""
        return f"{base}/chat/completions"

    def _build_headers(self) -> dict:
        headers = {"Content-Type": "application/json"}
        if self.provider.api_key:
            headers["Authorization"] = f"Bearer {self.provider.api_key}"
        return headers

    def _build_body(self, request: BridgeRequest, model: ModelConfig) -> dict:
        body: dict = {
            "model": model.id,
            "messages": request.messages,
            "stream": request.stream,
        }
        if request.temperature is not None:
            body["temperature"] = request.temperature
        if request.max_tokens is not None:
            body["max_tokens"] = request.max_tokens
        if request.top_p is not None:
            body["top_p"] = request.top_p
        if request.tools:
            body["tools"] = request.tools
        if request.tool_choice is not None:
            body["tool_choice"] = request.tool_choice
        if request.stop:
            body["stop"] = request.stop
        # Pass through any extra fields
        for k, v in request.extra.items():
            if k not in body:
                body[k] = v
        return body

    async def complete(self, request: BridgeRequest, model: ModelConfig) -> BridgeResponse:
        url = self._build_url()
        headers = self._build_headers()
        body = self._build_body(request, model)
        body["stream"] = False

        async with self.session.post(url, json=body, headers=headers) as resp:
            if resp.status != 200:
                error_text = await resp.text()
                logger.error(f"OpenAI API error {resp.status}: {error_text[:500]}")
                raise aiohttp.ClientResponseError(
                    resp.request_info, resp.history,
                    status=resp.status, message=error_text[:200],
                )

            data = await resp.json()
            choice = data.get("choices", [{}])[0]
            message = choice.get("message", {})
            usage = data.get("usage", {})

            return BridgeResponse(
                content=message.get("content", ""),
                model=data.get("model", model.id),
                finish_reason=choice.get("finish_reason"),
                tokens_in=usage.get("prompt_tokens", 0),
                tokens_out=usage.get("completion_tokens", 0),
                tool_calls=message.get("tool_calls"),
                raw=data,
            )

    async def stream(self, request: BridgeRequest, model: ModelConfig) -> StreamResult:
        """Stream OpenAI SSE with usage capture.
        
        Requests usage in the final chunk via stream_options.
        """
        result = StreamResult.__new__(StreamResult)
        result.usage = StreamUsage()
        result._iterator = self._stream_inner(request, model, result.usage)
        return result

    async def _stream_inner(
        self, request: BridgeRequest, model: ModelConfig, usage: StreamUsage
    ) -> AsyncIterator[str]:
        url = self._build_url()
        headers = self._build_headers()
        body = self._build_body(request, model)
        body["stream"] = True
        # Request usage in the final chunk (OpenAI feature)
        body["stream_options"] = {"include_usage": True}

        async with self.session.post(url, json=body, headers=headers) as resp:
            if resp.status != 200:
                error_text = await resp.text()
                logger.error(f"OpenAI API error {resp.status}: {error_text[:500]}")
                raise aiohttp.ClientResponseError(
                    resp.request_info, resp.history,
                    status=resp.status, message=error_text[:200],
                )

            async for line in resp.content:
                decoded = line.decode("utf-8", errors="replace").strip()
                if not decoded:
                    continue
                    
                # Parse usage from the final chunk
                if decoded.startswith("data: ") and decoded != "data: [DONE]":
                    try:
                        data = json.loads(decoded[6:])
                        chunk_usage = data.get("usage")
                        if chunk_usage:
                            usage.tokens_in = chunk_usage.get("prompt_tokens", 0)
                            usage.tokens_out = chunk_usage.get("completion_tokens", 0)
                    except json.JSONDecodeError:
                        pass
                
                yield decoded
