"""OpenAI client implementation."""

import os
import logging
from typing import Dict, List, Optional, Any
import aiohttp
import base64
import re
import json
import ssl
import certifi
from datetime import datetime
from .base import BaseOmniClient

logger = logging.getLogger(__name__)

# OpenAI specific client for the OmniLoop


class OpenAIClient(BaseOmniClient):
    """OpenAI vision API client implementation."""

    def __init__(
        self,
        api_key: Optional[str] = None,
        model: str = "gpt-4o",
        provider_base_url: str = "https://api.openai.com/v1",
        max_tokens: int = 4096,
        temperature: float = 0.0,
    ):
        """Initialize the OpenAI client.

        Args:
            api_key: OpenAI API key
            model: Model to use
            provider_base_url: API endpoint
            max_tokens: Maximum tokens to generate
            temperature: Generation temperature
        """
        super().__init__(api_key=api_key, model=model)
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        if not self.api_key:
            raise ValueError("No OpenAI API key provided")

        self.model = model
        self.provider_base_url = provider_base_url
        self.max_tokens = max_tokens
        self.temperature = temperature

    def _extract_base64_image(self, text: str) -> Optional[str]:
        """Extract base64 image data from an HTML img tag."""
        pattern = r'data:image/[^;]+;base64,([^"]+)'
        match = re.search(pattern, text)
        return match.group(1) if match else None

    def _get_loggable_messages(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Create a loggable version of messages with image data truncated."""
        loggable_messages = []
        for msg in messages:
            if isinstance(msg.get("content"), list):
                new_content = []
                for content in msg["content"]:
                    if content.get("type") == "image":
                        new_content.append(
                            {"type": "image", "image_url": {"url": "[BASE64_IMAGE_DATA]"}}
                        )
                    else:
                        new_content.append(content)
                loggable_messages.append({"role": msg["role"], "content": new_content})
            else:
                loggable_messages.append(msg)
        return loggable_messages

    async def run_interleaved(
        self, messages: List[Dict[str, Any]], system: str, max_tokens: Optional[int] = None
    ) -> Dict[str, Any]:
        """Run interleaved chat completion.

        Args:
            messages: List of message dicts
            system: System prompt
            max_tokens: Optional max tokens override

        Returns:
            Response dict
        """
        headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"}

        final_messages = [{"role": "system", "content": system}]

        # Process messages
        for item in messages:
            if isinstance(item, dict):
                if isinstance(item["content"], list):
                    # Content is already in the correct format
                    final_messages.append(item)
                else:
                    # Single string content, check for image
                    base64_img = self._extract_base64_image(item["content"])
                    if base64_img:
                        message = {
                            "role": item["role"],
                            "content": [
                                {
                                    "type": "image_url",
                                    "image_url": {"url": f"data:image/jpeg;base64,{base64_img}"},
                                }
                            ],
                        }
                    else:
                        message = {
                            "role": item["role"],
                            "content": [{"type": "text", "text": item["content"]}],
                        }
                    final_messages.append(message)
            else:
                # String content, check for image
                base64_img = self._extract_base64_image(item)
                if base64_img:
                    message = {
                        "role": "user",
                        "content": [
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/jpeg;base64,{base64_img}"},
                            }
                        ],
                    }
                else:
                    message = {"role": "user", "content": [{"type": "text", "text": item}]}
                final_messages.append(message)

        payload = {"model": self.model, "messages": final_messages, "temperature": self.temperature}

        if "o1" in self.model or "o3-mini" in self.model:
            payload["reasoning_effort"] = "low"
            payload["max_completion_tokens"] = max_tokens or self.max_tokens
        else:
            payload["max_tokens"] = max_tokens or self.max_tokens

        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    f"{self.provider_base_url}/chat/completions", headers=headers, json=payload
                ) as response:
                    response_json = await response.json()

                    if response.status != 200:
                        error_msg = response_json.get("error", {}).get(
                            "message", str(response_json)
                        )
                        logger.error(f"Error in OpenAI API call: {error_msg}")
                        raise Exception(f"OpenAI API error: {error_msg}")

                    return response_json

        except Exception as e:
            logger.error(f"Error in OpenAI API call: {str(e)}")
            raise
