"""Thin wrapper around the OpenAI Responses API."""

from __future__ import annotations

import base64
import mimetypes
from pathlib import Path
import logging
from typing import Optional, Sequence

from openai import OpenAI

from .types import ImageInput, LLMResponse, MessageSequence, TokenUsage, normalize_messages
from .utils import clamp_retries, run_sync_in_thread, run_with_retries

logger = logging.getLogger(__name__)


class OpenAIResponsesClient:
    """Convenience wrapper around the OpenAI Responses API."""

    def generate_response(
        self,
        *,
        api_key: str,
        prompt: Optional[str] = None,
        model: str,
        max_tokens: int = 32000,
        reasoning_effort: Optional[str] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        images: Optional[Sequence[ImageInput]] = None,
        messages: Optional[MessageSequence] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> str:
        """Generate a response from the specified model.

        Args:
            api_key: Secret key used to authenticate with OpenAI.
            prompt: Natural-language instruction or query for the model.
            model: Identifier of the OpenAI model to target (for example, ``"gpt-4o"``).
            max_tokens: Cap for tokens across the entire exchange, defaults to 32000.
            reasoning_effort: Optional reasoning effort hint (``"low"``, ``"medium"``, or ``"high"``).
            temperature: Optional sampling temperature.
            top_p: Optional nucleus sampling value.
            images: Optional collection of image references (local paths or URLs).
            messages: Optional list of chat-style messages (role/content).
            request_id: Optional request identifier for tracing/logging.
            timeout_s: Optional request timeout in seconds.
            max_retries: Optional retry count for transient failures.
            retry_backoff_s: Base delay (seconds) for exponential backoff between retries.

        Returns:
            The text output produced by the model.

        Raises:
            ValueError: If required arguments are missing or the request payload is empty.
            OpenAIError: If the underlying OpenAI request fails.
        """
        if not api_key:
            raise ValueError("api_key must be provided.")
        if not prompt and not messages and not images:
            raise ValueError("At least one of prompt, messages, or images must be provided.")
        if not model:
            raise ValueError("model must be provided.")

        normalized_messages = normalize_messages(prompt=prompt, messages=messages)
        content_messages: list[dict] = []
        for message in normalized_messages:
            content_blocks = []
            if message["content"]:
                content_blocks.append({"type": "input_text", "text": message["content"]})
            content_messages.append({"role": message["role"], "content": content_blocks})

        if images:
            image_blocks = [self._to_image_block(image) for image in images]
            target_index = next(
                (
                    index
                    for index in range(len(content_messages) - 1, -1, -1)
                    if content_messages[index]["role"] == "user"
                ),
                None,
            )
            if target_index is None:
                content_messages.append({"role": "user", "content": image_blocks})
            else:
                content_messages[target_index]["content"].extend(image_blocks)

        if not any(message["content"] for message in content_messages):
            raise ValueError("No content provided for response generation.")

        request_payload = {
            "model": model,
            "input": content_messages,
            "max_output_tokens": max_tokens,
        }

        if reasoning_effort:
            request_payload["reasoning"] = {"effort": reasoning_effort}
        if temperature is not None:
            request_payload["temperature"] = float(temperature)
        if top_p is not None:
            request_payload["top_p"] = float(top_p)

        retry_count = clamp_retries(max_retries)

        def _run_request() -> str:
            client_kwargs = {"api_key": api_key}
            if timeout_s is not None:
                client_kwargs["timeout"] = timeout_s
            client = OpenAI(**client_kwargs)
            try:
                response = client.responses.create(**request_payload)
            except Exception as exc:  # Log and re-raise to preserve default behavior
                logger.exception(
                    "OpenAI Responses API request failed: %s request_id=%s",
                    exc,
                    request_id,
                )
                raise

            output_text = response.output_text
            logger.info(
                "OpenAI generate_response succeeded: model=%s images=%d text_len=%d request_id=%s",
                model,
                len(images or []),
                len(output_text or ""),
                request_id,
            )
            return output_text

        return run_with_retries(
            func=_run_request,
            max_retries=retry_count,
            retry_backoff_s=retry_backoff_s,
            request_id=request_id,
        )

    def generate_response_with_usage(
        self,
        *,
        api_key: str,
        prompt: Optional[str] = None,
        model: str,
        max_tokens: int = 32000,
        reasoning_effort: Optional[str] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        images: Optional[Sequence[ImageInput]] = None,
        messages: Optional[MessageSequence] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> LLMResponse:
        if not api_key:
            raise ValueError("api_key must be provided.")
        if not prompt and not messages and not images:
            raise ValueError("At least one of prompt, messages, or images must be provided.")
        if not model:
            raise ValueError("model must be provided.")

        normalized_messages = normalize_messages(prompt=prompt, messages=messages)
        content_messages: list[dict] = []
        for message in normalized_messages:
            content_blocks = []
            if message["content"]:
                content_blocks.append({"type": "input_text", "text": message["content"]})
            content_messages.append({"role": message["role"], "content": content_blocks})

        if images:
            image_blocks = [self._to_image_block(image) for image in images]
            target_index = next(
                (
                    index
                    for index in range(len(content_messages) - 1, -1, -1)
                    if content_messages[index]["role"] == "user"
                ),
                None,
            )
            if target_index is None:
                content_messages.append({"role": "user", "content": image_blocks})
            else:
                content_messages[target_index]["content"].extend(image_blocks)

        if not any(message["content"] for message in content_messages):
            raise ValueError("No content provided for response generation.")

        request_payload = {
            "model": model,
            "input": content_messages,
            "max_output_tokens": max_tokens,
        }

        if reasoning_effort:
            request_payload["reasoning"] = {"effort": reasoning_effort}
        if temperature is not None:
            request_payload["temperature"] = float(temperature)
        if top_p is not None:
            request_payload["top_p"] = float(top_p)

        retry_count = clamp_retries(max_retries)

        def _run_request() -> LLMResponse:
            client_kwargs = {"api_key": api_key}
            if timeout_s is not None:
                client_kwargs["timeout"] = timeout_s
            client = OpenAI(**client_kwargs)
            try:
                response = client.responses.create(**request_payload)
            except Exception as exc:  # Log and re-raise to preserve default behavior
                logger.exception(
                    "OpenAI Responses API request failed: %s request_id=%s",
                    exc,
                    request_id,
                )
                raise

            output_text = response.output_text
            usage = _extract_openai_usage(response)
            logger.info(
                "OpenAI generate_response_with_usage succeeded: model=%s images=%d text_len=%d request_id=%s",
                model,
                len(images or []),
                len(output_text or ""),
                request_id,
            )
            return LLMResponse(text=output_text, usage=usage, provider="openai", model=model)

        return run_with_retries(
            func=_run_request,
            max_retries=retry_count,
            retry_backoff_s=retry_backoff_s,
            request_id=request_id,
        )

    async def async_generate_response(
        self,
        *,
        api_key: str,
        prompt: Optional[str] = None,
        model: str,
        max_tokens: int = 32000,
        reasoning_effort: Optional[str] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        images: Optional[Sequence[ImageInput]] = None,
        messages: Optional[MessageSequence] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> str:
        return await run_sync_in_thread(
            lambda: self.generate_response(
                api_key=api_key,
                prompt=prompt,
                model=model,
                max_tokens=max_tokens,
                reasoning_effort=reasoning_effort,
                temperature=temperature,
                top_p=top_p,
                images=images,
                messages=messages,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )
        )

    @staticmethod
    def _to_image_block(image: ImageInput) -> dict:
        """Convert an image reference into an OpenAI Responses API block."""
        if isinstance(image, Path):
            return {
                "type": "input_image",
                "image_url": _encode_image_path(image),
            }

        if image.startswith(("http://", "https://", "data:")):
            return {
                "type": "input_image",
                "image_url": image,
            }

        return {
            "type": "input_image",
            "image_url": _encode_image_path(Path(image)),
        }

    def generate_image(
        self,
        *,
        api_key: str,
        prompt: str,
        model: str,
        image_size: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
        image: Optional[ImageInput] = None,
    ) -> bytes:
        """Generate an image using the OpenAI API.

        Raises:
            NotImplementedError: This method is not yet implemented for OpenAI.
        """
        raise NotImplementedError("Image generation is not implemented for OpenAI.")

    async def async_generate_image(
        self,
        *,
        api_key: str,
        prompt: str,
        model: str,
        image_size: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
        image: Optional[ImageInput] = None,
    ) -> bytes:
        return await run_sync_in_thread(
            lambda: self.generate_image(
                api_key=api_key,
                prompt=prompt,
                model=model,
                image_size=image_size,
                aspect_ratio=aspect_ratio,
                image=image,
            )
        )

    async def async_generate_response_with_usage(
        self,
        *,
        api_key: str,
        prompt: Optional[str] = None,
        model: str,
        max_tokens: int = 32000,
        reasoning_effort: Optional[str] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        images: Optional[Sequence[ImageInput]] = None,
        messages: Optional[MessageSequence] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> LLMResponse:
        return await run_sync_in_thread(
            lambda: self.generate_response_with_usage(
                api_key=api_key,
                prompt=prompt,
                model=model,
                max_tokens=max_tokens,
                reasoning_effort=reasoning_effort,
                temperature=temperature,
                top_p=top_p,
                images=images,
                messages=messages,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )
        )

    def list_models(
        self,
        *,
        api_key: str,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> list[dict[str, Optional[str]]]:
        """Return the models available to the authenticated OpenAI account."""
        if not api_key:
            raise ValueError("api_key must be provided.")

        retry_count = clamp_retries(max_retries)

        def _run_request() -> list[dict[str, Optional[str]]]:
            client_kwargs = {"api_key": api_key}
            if timeout_s is not None:
                client_kwargs["timeout"] = timeout_s
            client = OpenAI(**client_kwargs)
            try:
                response = client.models.list()
            except Exception as exc:
                logger.exception(
                    "OpenAI list models failed: %s request_id=%s", exc, request_id
                )
                raise
            data = getattr(response, "data", []) or []

            models: list[dict[str, Optional[str]]] = []
            for model in data:
                model_id = getattr(model, "id", None)
                if model_id is None and isinstance(model, dict):
                    model_id = model.get("id")
                if not model_id:
                    continue

                display_name = getattr(model, "display_name", None)
                if display_name is None and isinstance(model, dict):
                    display_name = model.get("display_name")

                models.append({"id": model_id, "display_name": display_name})

            logger.info(
                "OpenAI list_models succeeded: count=%d request_id=%s",
                len(models),
                request_id,
            )
            return models

        return run_with_retries(
            func=_run_request,
            max_retries=retry_count,
            retry_backoff_s=retry_backoff_s,
            request_id=request_id,
        )

    async def async_list_models(
        self,
        *,
        api_key: str,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> list[dict[str, Optional[str]]]:
        return await run_sync_in_thread(
            lambda: self.list_models(
                api_key=api_key,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )
        )


def _encode_image_path(path: Path) -> str:
    """Generate a data URL for the provided image path."""
    data = path.read_bytes()
    encoded = base64.b64encode(data).decode("utf-8")
    mime_type = mimetypes.guess_type(path.name)[0] or "image/png"
    return f"data:{mime_type};base64,{encoded}"


def _extract_openai_usage(response: object) -> TokenUsage | None:
    usage_obj = getattr(response, "usage", None)
    if usage_obj is None:
        return None

    input_tokens = getattr(usage_obj, "input_tokens", None)
    output_tokens = getattr(usage_obj, "output_tokens", None)
    total_tokens = getattr(usage_obj, "total_tokens", None)

    if isinstance(usage_obj, dict):
        input_tokens = usage_obj.get("input_tokens")
        output_tokens = usage_obj.get("output_tokens")
        total_tokens = usage_obj.get("total_tokens")

    return TokenUsage(
        input_tokens=int(input_tokens) if isinstance(input_tokens, int) else None,
        output_tokens=int(output_tokens) if isinstance(output_tokens, int) else None,
        total_tokens=int(total_tokens) if isinstance(total_tokens, int) else None,
    )
