"""OpenAI adapter: data models, conversion helpers, and async API wrapper."""

import base64
import json
import threading
from enum import StrEnum
from typing import Any, Literal

from openai import AsyncOpenAI, BadRequestError
from pydantic import BaseModel, ConfigDict, Field

from dotpromptz.image import resolve_image_ref, to_data_url
from dotpromptz.typing import (
    DataPart,
    MediaPart,
    Message,
    Part,
    RenderedPrompt,
    Role,
    TextPart,
)

from ._base import (
    Adapter,
    Credential,
    GeneratedImage,
    GenerateResponse,
    Usage,
    resolve_config_params,
    resolve_rendered_config_and_model,
)


class DetailKind(StrEnum):
    """The kind of Image URL detail."""

    AUTO = 'auto'
    LOW = 'low'
    HIGH = 'high'


# Define TypedDicts to represent the interfaces
class ImageURLDetail(BaseModel):
    """Image URL Detail."""

    url: str
    detail: DetailKind | None = None


class ContentItemType(StrEnum):
    """Enumveration variants for content item type."""

    TEXT = 'text'
    IMAGE_URL = 'image_url'


class ContentItem(BaseModel):
    """Content Item: can be text or image."""

    type: ContentItemType
    text: str | None = None
    image_url: ImageURLDetail | None = None


class OpenAIMessage(BaseModel):
    """Open AI Message."""

    role: str
    content: str | list[ContentItem] | None = None
    name: str | None = None


class ResponseFormatType(StrEnum):
    """Enum variants for response format type."""

    TEXT = 'text'
    JSON_OBJECT = 'json_object'
    JSON_SCHEMA = 'json_schema'


class JsonSchemaSpec(BaseModel):
    """JSON Schema specification within a structured output response format."""

    name: str
    strict: bool = True
    schema_value: dict[str, Any] = Field(
        validation_alias='schema',
        serialization_alias='schema',
    )
    model_config = ConfigDict(
        populate_by_name=True,
        serialize_by_alias=True,
    )


class ResponseFormat(BaseModel):
    """Expected Response Format."""

    type: ResponseFormatType
    json_schema: JsonSchemaSpec | None = None


ReasoningEffort = Literal['low', 'medium', 'high']


class OpenAIRequest(BaseModel):
    """Open AI Request."""

    messages: list[OpenAIMessage]
    model: str
    frequency_penalty: float | None = None
    logit_bias: dict[str, int] | None = None
    max_completion_tokens: int | None = None
    n: int | None = None
    presence_penalty: float | None = None
    reasoning_effort: ReasoningEffort | None = None
    response_format: ResponseFormat | None = None
    seed: int | None = None
    stop: str | list[str] | None = None
    temperature: float | None = None
    top_p: float | None = None
    user: str | None = None


# ---------------------------------------------------------------------------
# Conversion helpers
# ---------------------------------------------------------------------------

_ROLE_MAP: dict[Role, str] = {
    Role.USER: 'user',
    Role.MODEL: 'assistant',
    Role.SYSTEM: 'system',
}

_OPENAI_IMAGE_SIZE_MAP = {
    '1:1': '1024x1024',
    '3:2': '1536x1024',
    '2:3': '1024x1536',
}
_OPENAI_IMAGE_SIZES = frozenset({'auto', '1024x1024', '1536x1024', '1024x1536'})


def _role_to_openai(role: Role) -> str:
    """Map a dotprompt ``Role`` to the OpenAI role string.

    Args:
        role: The dotprompt role enum value.

    Returns:
        The corresponding OpenAI role string.
    """
    return _ROLE_MAP.get(role, role.value)


def _part_to_content_item(part: Part) -> ContentItem | None:
    """Convert a single dotprompt ``Part`` to an OpenAI content item.

    Args:
        part: A dotprompt content part.

    Returns:
        A ``ContentItem`` or ``None`` if unsupported.
    """
    if isinstance(part, TextPart):
        return ContentItem(type=ContentItemType.TEXT, text=part.text)
    if isinstance(part, MediaPart):
        resolved = resolve_image_ref(
            part.media.url,
            content_type=part.media.content_type,
            read_bytes=True,
        )
        image_url = resolved.url
        if resolved.kind == 'local_file':
            if resolved.data is None:
                raise ValueError('Local image bytes are missing during OpenAI request conversion.')
            image_url = to_data_url(resolved.data, resolved.mime_type)
        return ContentItem(
            type=ContentItemType.IMAGE_URL,
            image_url=ImageURLDetail(url=image_url or part.media.url),
        )
    if isinstance(part, DataPart):
        return ContentItem(type=ContentItemType.TEXT, text=json.dumps(part.data, ensure_ascii=False))
    return None


def _parts_to_content(parts: list[Part]) -> str | list[ContentItem]:
    """Convert a list of dotprompt ``Part`` objects to OpenAI content.

    If all parts are text-only, returns a single concatenated string.
    Otherwise returns a list of ``ContentItem``.

    Args:
        parts: List of dotprompt content parts.

    Returns:
        Either a plain string or a list of ``ContentItem``.
    """
    if not parts:
        return ''
    if all(isinstance(p, TextPart) for p in parts):
        return ''.join(p.text for p in parts if isinstance(p, TextPart))

    items: list[ContentItem] = []
    for part in parts:
        item = _part_to_content_item(part)
        if isinstance(item, ContentItem):
            items.append(item)
    return items


def _message_to_openai(msg: Message) -> OpenAIMessage:
    """Convert a single dotprompt ``Message`` to an ``OpenAIMessage``.

    Args:
        msg: A dotprompt message.

    Returns:
        An ``OpenAIMessage``.
    """
    role = _role_to_openai(msg.role)
    content = _parts_to_content(msg.content)

    return OpenAIMessage(
        role=role,
        content=content if content else None,
    )


def to_openai_messages(messages: list[Message]) -> list[OpenAIMessage]:
    """Convert dotprompt messages to OpenAI message format.

    Args:
        messages: List of dotprompt ``Message`` objects.

    Returns:
        List of ``OpenAIMessage`` objects.
    """
    return [_message_to_openai(m) for m in messages]


def to_openai_request(rendered: RenderedPrompt) -> OpenAIRequest:
    """Convert a ``RenderedPrompt`` to a full ``OpenAIRequest``.

    Args:
        rendered: The rendered prompt from ``Dotprompt.render()``.

    Returns:
        An ``OpenAIRequest`` ready for the OpenAI API.

    Raises:
        ValueError: If no model is specified.
    """
    config, resolved_model = resolve_rendered_config_and_model(rendered)

    response_format: ResponseFormat | None = None
    if rendered.output and rendered.output.format in ('json', 'yaml'):
        if rendered.output.schema_value is not None:
            response_format = ResponseFormat(
                type=ResponseFormatType.JSON_SCHEMA,
                json_schema=JsonSchemaSpec(
                    name='response',
                    strict=True,
                    schema_value=rendered.output.schema_value,
                ),
            )
        else:
            response_format = ResponseFormat(type=ResponseFormatType.JSON_OBJECT)

    params = resolve_config_params(config)

    return OpenAIRequest(
        messages=to_openai_messages(rendered.messages),
        model=resolved_model,
        temperature=params.get('temperature'),
        top_p=params.get('top_p'),
        max_completion_tokens=params.get('max_tokens'),
        stop=params.get('stop'),
        frequency_penalty=params.get('frequency_penalty'),
        presence_penalty=params.get('presence_penalty'),
        reasoning_effort=params.get('thinking'),
        seed=params.get('seed'),
        response_format=response_format,
    )


def _is_image_output(rendered: RenderedPrompt) -> bool:
    """Return whether this prompt requests image output."""
    return rendered.output is not None and rendered.output.format == 'image'


def _normalize_image_prompt_text(messages: list[Message]) -> str:
    """Build one image-generation prompt string from conversational messages."""
    role_labels = {
        Role.SYSTEM: 'System',
        Role.USER: 'User',
        Role.MODEL: 'Assistant',
    }

    lines: list[str] = []
    for msg in messages:
        text_parts: list[str] = []
        for part in msg.content:
            if isinstance(part, TextPart):
                text_parts.append(part.text)
            elif isinstance(part, DataPart):
                text_parts.append(json.dumps(part.data, ensure_ascii=False))
        text = '\n'.join(chunk for chunk in text_parts if chunk).strip()
        if text:
            lines.append(f'{role_labels[msg.role]}: {text}')

    prompt = '\n\n'.join(lines).strip()
    if not prompt:
        raise ValueError('Image generation requires at least one text instruction in prompt messages.')
    return prompt


def _resolve_openai_image_size(config: dict[str, Any]) -> str | None:
    """Resolve OpenAI image size from ``resolution``/``aspect_ratio`` config."""
    size_from_resolution: str | None = None
    raw_resolution = config.get('resolution')
    if raw_resolution is not None:
        if not isinstance(raw_resolution, str):
            raise ValueError('config.resolution must be a string for OpenAI image generation.')
        normalized_resolution = raw_resolution.strip().lower()
        if normalized_resolution not in _OPENAI_IMAGE_SIZES:
            raise ValueError(
                'config.resolution for OpenAI image generation must be one of '
                "['auto', '1024x1024', '1536x1024', '1024x1536']."
            )
        size_from_resolution = normalized_resolution

    size_from_aspect: str | None = None
    raw_aspect_ratio = config.get('aspect_ratio')
    if raw_aspect_ratio is not None:
        if not isinstance(raw_aspect_ratio, str):
            raise ValueError('config.aspect_ratio must be a string for OpenAI image generation.')
        normalized_aspect_ratio = raw_aspect_ratio.replace(' ', '')
        if normalized_aspect_ratio not in _OPENAI_IMAGE_SIZE_MAP:
            raise ValueError("config.aspect_ratio for OpenAI image generation must be one of ['1:1', '3:2', '2:3'].")
        size_from_aspect = _OPENAI_IMAGE_SIZE_MAP[normalized_aspect_ratio]

    if size_from_resolution and size_from_aspect and size_from_resolution != size_from_aspect:
        raise ValueError(
            'config.resolution and config.aspect_ratio conflict for OpenAI image generation: '
            f'{size_from_resolution!r} vs {size_from_aspect!r}.'
        )
    return size_from_resolution or size_from_aspect


def _resolve_openai_background(config: dict[str, Any]) -> str | None:
    """Resolve OpenAI background mode from ``transparent`` config."""
    raw_transparent = config.get('transparent')
    if raw_transparent is None:
        return None
    if not isinstance(raw_transparent, bool):
        raise ValueError('config.transparent must be a boolean for OpenAI image generation.')
    return 'transparent' if raw_transparent else 'opaque'


def _mime_to_suffix(mime_type: str) -> str:
    """Map MIME type to a filename suffix for multipart uploads."""
    suffix_map = {
        'image/png': '.png',
        'image/jpeg': '.jpg',
        'image/webp': '.webp',
    }
    return suffix_map.get(mime_type.lower(), '.png')


def _collect_openai_edit_images(messages: list[Message]) -> list[tuple[str, bytes, str]]:
    """Collect media parts as multipart image inputs for OpenAI edit API."""
    images: list[tuple[str, bytes, str]] = []
    for idx, msg in enumerate(messages):
        for part_index, part in enumerate(msg.content):
            if not isinstance(part, MediaPart):
                continue
            resolved = resolve_image_ref(
                part.media.url,
                content_type=part.media.content_type,
                read_bytes=True,
                read_remote_bytes=True,
            )
            if resolved.data is None:
                raise ValueError('Image bytes are missing for OpenAI image edit input.')
            suffix = _mime_to_suffix(resolved.mime_type)
            name = f'input_{idx}_{part_index}{suffix}'
            images.append((name, resolved.data, resolved.mime_type))
    return images


def to_openai_image_request(rendered: RenderedPrompt) -> dict[str, Any]:
    """Convert a ``RenderedPrompt`` to OpenAI image request parameters."""
    config, resolved_model = resolve_rendered_config_and_model(rendered)
    request: dict[str, Any] = {
        'model': resolved_model,
        'prompt': _normalize_image_prompt_text(rendered.messages),
        'output_format': 'png',
    }

    size = _resolve_openai_image_size(config)
    if size is not None:
        request['size'] = size

    background = _resolve_openai_background(config)
    if background is not None:
        request['background'] = background

    return request


def _extract_openai_generated_images(response: Any) -> list[GeneratedImage]:
    """Extract generated image bytes from OpenAI image response payload."""
    output: list[GeneratedImage] = []
    data_items = getattr(response, 'data', None)
    if not isinstance(data_items, list):
        raise ValueError('OpenAI image response did not include image data.')

    for item in data_items:
        b64_json = getattr(item, 'b64_json', None)
        if not isinstance(b64_json, str) or not b64_json:
            continue
        output.append(
            GeneratedImage(
                data=base64.b64decode(b64_json),
                mime_type='image/png',
            )
        )

    if not output:
        raise ValueError('OpenAI image response contained no decodable images.')
    return output


def _build_usage_from_openai_image_response(response: Any) -> Usage | None:
    """Build unified usage from an OpenAI image response."""
    usage = getattr(response, 'usage', None)
    if usage is None:
        return None
    prompt_tokens = getattr(usage, 'input_tokens', 0) or 0
    completion_tokens = getattr(usage, 'output_tokens', 0) or 0
    total_tokens = getattr(usage, 'total_tokens', None)
    if total_tokens is None:
        total_tokens = prompt_tokens + completion_tokens
    return Usage(
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=total_tokens,
    )


def _extract_text_from_openai_content(content: Any) -> str | None:
    """Extract plain text from OpenAI message content payloads."""
    if isinstance(content, str):
        return content
    if not isinstance(content, list):
        return None

    chunks: list[str] = []
    for part in content:
        if isinstance(part, str):
            chunks.append(part)
            continue
        if isinstance(part, dict):
            raw_text = part.get('text')
            if isinstance(raw_text, str):
                chunks.append(raw_text)
                continue
            if isinstance(raw_text, dict):
                value = raw_text.get('value')
                if isinstance(value, str):
                    chunks.append(value)
            continue

        raw_text = getattr(part, 'text', None)
        if isinstance(raw_text, str):
            chunks.append(raw_text)
            continue
        if isinstance(raw_text, dict):
            value = raw_text.get('value')
            if isinstance(value, str):
                chunks.append(value)
            continue
        value = getattr(raw_text, 'value', None)
        if isinstance(value, str):
            chunks.append(value)

    if not chunks:
        return None
    return ''.join(chunks)


# ---------------------------------------------------------------------------
# OpenAI Adapter
# ---------------------------------------------------------------------------


class OpenAIAdapter(Adapter):
    """Adapter for the OpenAI Chat Completions API.

    Args:
        credential: A ``Credential`` with ``api_key`` and optional ``base_url``.
        default_model: Fallback model when the prompt has none.
        organization: Optional OpenAI organization ID.
    """

    def __init__(
        self,
        *,
        credential: Credential,
        default_model: str = 'gpt-5.2',
        organization: str | None = None,
    ) -> None:
        """Initialise the OpenAI adapter."""
        self._credential = credential
        self._default_model = default_model
        self._organization = organization
        self._client: Any = None
        self._client_lock = threading.Lock()

    def _get_client(self) -> Any:
        """Return (and lazily create) the ``AsyncOpenAI`` client."""

        def _factory() -> AsyncOpenAI:
            kwargs: dict[str, Any] = {'api_key': self._credential.api_key}
            if self._credential.base_url:
                kwargs['base_url'] = self._credential.base_url
            if self._organization:
                kwargs['organization'] = self._organization
            return AsyncOpenAI(**kwargs)

        return self._get_or_create_client(_factory)

    def convert(self, rendered: RenderedPrompt) -> dict[str, Any]:
        """Convert a ``RenderedPrompt`` to an OpenAI request dict.

        Args:
            rendered: The rendered prompt.

        Returns:
            A dict suitable for ``client.chat.completions.create(**d)``.
        """
        prepared = self._ensure_model(rendered)
        if _is_image_output(prepared):
            return to_openai_image_request(prepared)
        req = to_openai_request(prepared)
        return req.model_dump(exclude_none=True)

    async def generate(
        self,
        rendered: RenderedPrompt,
        **kwargs: Any,
    ) -> GenerateResponse:
        """Call the OpenAI Chat Completions API.

        Args:
            rendered: The rendered prompt from ``Dotprompt.render()``.
            **kwargs: Extra keyword args forwarded to the SDK.

        Returns:
            A ``GenerateResponse``.
        """
        if _is_image_output(rendered):
            return await self._generate_image(rendered, **kwargs)

        client = self._get_client()
        request_dict = self.convert(rendered)
        request_dict.update(kwargs)

        is_structured_output = rendered.output and rendered.output.format in ('json', 'yaml')

        async def _request_once() -> Any:
            if not is_structured_output:
                return await client.chat.completions.create(**request_dict)
            if rendered.output and rendered.output.schema_value is None:
                return await client.chat.completions.create(**request_dict)
            return await client.chat.completions.parse(**request_dict)

        try:
            response = await _request_once()
        except BadRequestError as exc:
            message = str(exc).lower()
            if 'invalid image data' in message:
                raise ValueError(
                    'OpenAI endpoint rejected local image payload. '
                    'This endpoint may not support data-URL image inputs. '
                    'Use an http/https image URL instead.'
                ) from exc
            raise

        choice = response.choices[0]
        msg = choice.message
        text = _extract_text_from_openai_content(msg.content)
        parsed = getattr(msg, 'parsed', None)
        if parsed is not None:
            if isinstance(parsed, BaseModel):
                text = json.dumps(parsed.model_dump(mode='json'), ensure_ascii=False)
            else:
                text = json.dumps(parsed, ensure_ascii=False)

        if is_structured_output and not text and not msg.refusal:
            # Retry once for occasional empty content in structured-output mode.
            try:
                retry_response = await _request_once()
            except BadRequestError:
                retry_response = None
            if retry_response is not None:
                response = retry_response
                choice = response.choices[0]
                msg = choice.message
                text = _extract_text_from_openai_content(msg.content)
                parsed = getattr(msg, 'parsed', None)
                if parsed is not None:
                    if isinstance(parsed, BaseModel):
                        text = json.dumps(parsed.model_dump(mode='json'), ensure_ascii=False)
                    else:
                        text = json.dumps(parsed, ensure_ascii=False)

        usage: Usage | None = None
        if response.usage:
            usage = Usage(
                prompt_tokens=response.usage.prompt_tokens,
                completion_tokens=response.usage.completion_tokens,
                total_tokens=response.usage.total_tokens,
            )

        return GenerateResponse(
            text=text,
            usage=usage,
            raw=response,
            model=response.model,
            finish_reason=choice.finish_reason,
        )

    async def _generate_image(
        self,
        rendered: RenderedPrompt,
        **kwargs: Any,
    ) -> GenerateResponse:
        """Call the OpenAI image API for image generation or editing."""
        client = self._get_client()
        prepared = self._ensure_model(rendered)
        request_dict = to_openai_image_request(prepared)
        request_dict.update(kwargs)

        edit_images = _collect_openai_edit_images(prepared.messages)
        if edit_images:
            response = await client.images.edit(
                image=edit_images,
                **request_dict,
            )
        else:
            response = await client.images.generate(**request_dict)

        return GenerateResponse(
            text=None,
            usage=_build_usage_from_openai_image_response(response),
            raw=response,
            model=request_dict.get('model'),
            finish_reason='stop',
            images=_extract_openai_generated_images(response),
        )
