from __future__ import annotations

import os
from typing import Any

import openai
import tiktoken

from nexus.protocols.providers.base import AbstractProvider, LLMRequest, LLMResponse, ToolCall
from nexus.resilience.errors import ProviderError


class OpenAIProvider(AbstractProvider):
    """OpenAI SDK-backed provider."""

    def __init__(self, model: str = "gpt-4o", api_key: str | None = None, base_url: str | None = None) -> None:
        self._model = model
        self._api_key = api_key or os.getenv("OPENAI_API_KEY")
        self._client = openai.AsyncOpenAI(api_key=self._api_key, base_url=base_url)

    async def complete(self, request: LLMRequest) -> LLMResponse:
        """Run chat completion and normalize response."""
        messages: list[dict[str, Any]] = []
        if request.system:
            messages.append({"role": "system", "content": request.system})
        for m in request.messages:
            item: dict[str, Any] = {"role": m.role, "content": m.content}
            if m.name:
                item["name"] = m.name
            if m.tool_call_id:
                item["tool_call_id"] = m.tool_call_id
            messages.append(item)

        tools = [
            {
                "type": "function",
                "function": {
                    "name": t.name,
                    "description": t.description,
                    "parameters": t.parameters,
                },
            }
            for t in request.tools
        ]

        try:
            resp = await self._client.chat.completions.create(
                model=self._model,
                messages=messages,
                tools=tools or openai.NOT_GIVEN,
                temperature=request.temperature,
                max_tokens=request.max_tokens,
                stop=request.stop_sequences or openai.NOT_GIVEN,
            )
        except Exception as exc:
            raise ProviderError(f"OpenAI call failed: {exc}", provider_name="openai") from exc

        choice = resp.choices[0]
        msg = choice.message
        calls: list[ToolCall] = []
        for tc in msg.tool_calls or []:
            calls.append(
                ToolCall(
                    tool_call_id=tc.id,
                    tool_name=tc.function.name,
                    arguments=_safe_json_loads(tc.function.arguments),
                )
            )

        usage = resp.usage
        return LLMResponse(
            content=msg.content or "",
            tool_calls=calls,
            input_tokens=usage.prompt_tokens if usage else 0,
            output_tokens=usage.completion_tokens if usage else 0,
            model=resp.model,
            finish_reason=choice.finish_reason or "stop",
            raw=resp.model_dump(),
        )

    def count_tokens(self, text: str) -> int:
        """Estimate token count for model with cl100k fallback."""
        try:
            enc = tiktoken.encoding_for_model(self._model)
        except KeyError:
            enc = tiktoken.get_encoding("cl100k_base")
        return len(enc.encode(text))

    @property
    def model_name(self) -> str:
        return self._model

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

    @property
    def context_window(self) -> int:
        windows = {
            "gpt-4o": 128000,
            "gpt-4-turbo": 128000,
            "gpt-3.5-turbo": 16385,
            "o3": 200000,
            "o3-mini": 200000,
        }
        return windows.get(self._model, 8192)


def _safe_json_loads(raw: str) -> dict[str, Any]:
    import json

    try:
        return json.loads(raw) if raw else {}
    except json.JSONDecodeError:
        return {"_raw": raw}
