import asyncio
import json
import logging
import re
from typing import Any

import httpx

from solar_eval.config import get_provider_config
from solar_eval.providers.base import BaseProvider

logger = logging.getLogger(__name__)

MAX_RETRIES = 3


class OpenAIProvider(BaseProvider):
    def __init__(self) -> None:
        self.api_key = get_provider_config().openai_api_key

    async def chat_completion(
        self,
        messages: list[dict[str, str]],
        model: str = "gpt-4o",
        temperature: float = 0.0,
        max_tokens: int = 4096,
        reasoning_effort: str | None = None,
        response_format: dict[str, Any] | None = None,
        tools: list[dict[str, Any]] | None = None,
        tool_choice: str | dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        if response_format:
            payload["response_format"] = response_format
        if tools:
            payload["tools"] = tools
            if tool_choice is not None:
                payload["tool_choice"] = tool_choice

        for attempt in range(MAX_RETRIES):
            try:
                async with httpx.AsyncClient(timeout=120.0) as client:
                    response = await client.post(
                        "https://api.openai.com/v1/chat/completions",
                        headers={"Authorization": f"Bearer {self.api_key}"},
                        json=payload,
                    )
                    response.raise_for_status()
                    data = response.json()

                choice = data["choices"][0]
                message = choice["message"]
                return {
                    "content": message.get("content") or "",
                    "usage": data.get("usage", {}),
                    "finish_reason": choice.get("finish_reason", "stop"),
                    "tool_calls": message.get("tool_calls") or [],
                }
            except (httpx.HTTPStatusError, httpx.TimeoutException) as e:
                logger.warning(f"OpenAI API attempt {attempt + 1} failed: {e}")
                if attempt == MAX_RETRIES - 1:
                    raise
                await asyncio.sleep(2**attempt)

        raise RuntimeError("Max retries exceeded")

    async def chat_completion_json(
        self,
        messages: list[dict[str, str]],
        model: str = "gpt-4o",
        temperature: float = 0.0,
        max_tokens: int = 4096,
    ) -> dict[str, Any]:
        result = await self.chat_completion(
            messages=messages, model=model, temperature=temperature, max_tokens=max_tokens
        )
        content = result["content"]
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            json_match = re.search(r"```(?:json)?\s*(.*?)```", content, re.DOTALL)
            if json_match:
                return json.loads(json_match.group(1).strip())
            raise