# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
# SPDX-License-Identifier: Apache-2.0
"""OpenAI VLM backend implementation"""

import asyncio
import base64
import logging
from pathlib import Path
from typing import Any, Dict, List, Union

from ..base import VLMBase

logger = logging.getLogger(__name__)


class OpenAIVLM(VLMBase):
    """OpenAI VLM backend"""

    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        self._sync_client = None
        self._async_client = None
        self.provider = "openai"

    def get_client(self):
        """Get sync client"""
        if self._sync_client is None:
            try:
                import openai
            except ImportError:
                raise ImportError("Please install openai: pip install openai")
            client_kwargs = {"api_key": self.api_key, "base_url": self.api_base}
            if self.extra_headers:
                client_kwargs["default_headers"] = self.extra_headers
            self._sync_client = openai.OpenAI(**client_kwargs)
        return self._sync_client

    def get_async_client(self):
        """Get async client"""
        if self._async_client is None:
            try:
                import openai
            except ImportError:
                raise ImportError("Please install openai: pip install openai")
            client_kwargs = {"api_key": self.api_key, "base_url": self.api_base}
            if self.extra_headers:
                client_kwargs["default_headers"] = self.extra_headers
            self._async_client = openai.AsyncOpenAI(**client_kwargs)
        return self._async_client

    def _update_token_usage_from_response(self, response):
        if hasattr(response, "usage") and response.usage:
            prompt_tokens = response.usage.prompt_tokens
            completion_tokens = response.usage.completion_tokens
            self.update_token_usage(
                model_name=self.model or "gpt-4o-mini",
                provider=self.provider,
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
            )
        return

    def _extract_from_chunk(self, chunk):
        """Extract content and usage from a single chunk.

        Returns:
            tuple: (content, prompt_tokens, completion_tokens)
        """
        content = None
        prompt_tokens = 0
        completion_tokens = 0

        # Extract content from delta
        if chunk.choices and chunk.choices[0].delta:
            content = getattr(chunk.choices[0].delta, "content", None)

        # Extract usage from chunk if available
        if hasattr(chunk, "usage") and chunk.usage:
            prompt_tokens = chunk.usage.prompt_tokens or 0
            completion_tokens = chunk.usage.completion_tokens or 0

        return content, prompt_tokens, completion_tokens

    def _process_streaming_response(self, response):
        """Process streaming response and extract content and token usage.

        Args:
            response: Streaming response iterator from OpenAI client

        Returns:
            str: Extracted content
        """
        content_parts = []
        prompt_tokens = 0
        completion_tokens = 0

        for chunk in response:
            content, pt, ct = self._extract_from_chunk(chunk)
            if content:
                content_parts.append(content)
            if pt > 0:
                prompt_tokens = pt
            if ct > 0:
                completion_tokens = ct

        # Update token usage if we got it from streaming chunks
        if prompt_tokens > 0 or completion_tokens > 0:
            self.update_token_usage(
                model_name=self.model or "gpt-4o-mini",
                provider=self.provider,
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
            )

        return "".join(content_parts)

    async def _process_streaming_response_async(self, response):
        """Process async streaming response and extract content and token usage.

        Args:
            response: Async streaming response iterator from OpenAI client

        Returns:
            str: Extracted content
        """
        content_parts = []
        prompt_tokens = 0
        completion_tokens = 0

        async for chunk in response:
            content, pt, ct = self._extract_from_chunk(chunk)
            if content:
                content_parts.append(content)
            if pt > 0:
                prompt_tokens = pt
            if ct > 0:
                completion_tokens = ct

        # Update token usage if we got it from streaming chunks
        if prompt_tokens > 0 or completion_tokens > 0:
            self.update_token_usage(
                model_name=self.model or "gpt-4o-mini",
                provider=self.provider,
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
            )

        return "".join(content_parts)

    def get_completion(self, prompt: str, thinking: bool = False) -> str:
        """Get text completion"""
        client = self.get_client()
        kwargs = {
            "model": self.model or "gpt-4o-mini",
            "messages": [{"role": "user", "content": prompt}],
            "temperature": self.temperature,
            "stream": self.stream,
        }
        if self.max_tokens is not None:
            kwargs["max_tokens"] = self.max_tokens

        response = client.chat.completions.create(**kwargs)

        if self.stream:
            content = self._process_streaming_response(response)
        else:
            self._update_token_usage_from_response(response)
            content = response.choices[0].message.content or ""

        return self._clean_response(content)

    async def get_completion_async(
        self, prompt: str, thinking: bool = False, max_retries: int = 0
    ) -> str:
        """Get text completion asynchronously"""
        client = self.get_async_client()
        kwargs = {
            "model": self.model or "gpt-4o-mini",
            "messages": [{"role": "user", "content": prompt}],
            "temperature": self.temperature,
            "stream": self.stream,
        }
        if self.max_tokens is not None:
            kwargs["max_tokens"] = self.max_tokens

        last_error = None
        for attempt in range(max_retries + 1):
            try:
                response = await client.chat.completions.create(**kwargs)

                if self.stream:
                    content = await self._process_streaming_response_async(response)
                else:
                    self._update_token_usage_from_response(response)
                    content = response.choices[0].message.content or ""

                return self._clean_response(content)
            except Exception as e:
                last_error = e
                if attempt < max_retries:
                    await asyncio.sleep(2**attempt)

        if last_error:
            raise last_error
        else:
            raise RuntimeError("Unknown error in async completion")

    def _detect_image_format(self, data: bytes) -> str:
        """Detect image format from magic bytes.

        Supported formats: PNG, JPEG, GIF, WebP
        """
        if len(data) < 8:
            logger.warning(f"[OpenAIVLM] Image data too small: {len(data)} bytes")
            return "image/png"

        if data[:8] == b"\x89PNG\r\n\x1a\n":
            return "image/png"
        elif data[:2] == b"\xff\xd8":
            return "image/jpeg"
        elif data[:6] in (b"GIF87a", b"GIF89a"):
            return "image/gif"
        elif data[:4] == b"RIFF" and len(data) >= 12 and data[8:12] == b"WEBP":
            return "image/webp"

        logger.warning(f"[OpenAIVLM] Unknown image format, magic bytes: {data[:8].hex()}")
        return "image/png"

    def _prepare_image(self, image: Union[str, Path, bytes]) -> Dict[str, Any]:
        """Prepare image data for vision completion."""
        if isinstance(image, bytes):
            b64 = base64.b64encode(image).decode("utf-8")
            mime_type = self._detect_image_format(image)
            return {
                "type": "image_url",
                "image_url": {"url": f"data:{mime_type};base64,{b64}"},
            }
        elif isinstance(image, Path) or (
            isinstance(image, str) and not image.startswith(("http://", "https://"))
        ):
            path = Path(image)
            suffix = path.suffix.lower()
            mime_type = {
                ".png": "image/png",
                ".jpg": "image/jpeg",
                ".jpeg": "image/jpeg",
                ".gif": "image/gif",
                ".webp": "image/webp",
            }.get(suffix, "image/png")
            with open(path, "rb") as f:
                data = f.read()
            b64 = base64.b64encode(data).decode("utf-8")
            return {
                "type": "image_url",
                "image_url": {"url": f"data:{mime_type};base64,{b64}"},
            }
        else:
            return {"type": "image_url", "image_url": {"url": image}}

    def get_vision_completion(
        self,
        prompt: str,
        images: List[Union[str, Path, bytes]],
        thinking: bool = False,
    ) -> str:
        """Get vision completion"""
        client = self.get_client()

        content = []
        for img in images:
            content.append(self._prepare_image(img))
        content.append({"type": "text", "text": prompt})

        kwargs = {
            "model": self.model or "gpt-4o-mini",
            "messages": [{"role": "user", "content": content}],
            "temperature": self.temperature,
            "stream": self.stream,
        }
        if self.max_tokens is not None:
            kwargs["max_tokens"] = self.max_tokens

        response = client.chat.completions.create(**kwargs)

        if self.stream:
            content = self._process_streaming_response(response)
        else:
            self._update_token_usage_from_response(response)
            content = response.choices[0].message.content or ""

        return self._clean_response(content)

    async def get_vision_completion_async(
        self,
        prompt: str,
        images: List[Union[str, Path, bytes]],
        thinking: bool = False,
    ) -> str:
        """Get vision completion asynchronously"""
        client = self.get_async_client()

        content = []
        for img in images:
            content.append(self._prepare_image(img))
        content.append({"type": "text", "text": prompt})

        kwargs = {
            "model": self.model or "gpt-4o-mini",
            "messages": [{"role": "user", "content": content}],
            "temperature": self.temperature,
            "stream": self.stream,
        }
        if self.max_tokens is not None:
            kwargs["max_tokens"] = self.max_tokens

        response = await client.chat.completions.create(**kwargs)

        if self.stream:
            content = await self._process_streaming_response_async(response)
        else:
            self._update_token_usage_from_response(response)
            content = response.choices[0].message.content or ""

        return self._clean_response(content)
