"""Vision + Structured Output extraction via Azure OpenAI / OpenAI.

Sends images (URLs or base64 data URIs) alongside a text instruction to the
Chat Completions API with ``response_format: json_schema`` and returns the
result as a Pydantic model.

Both Azure OpenAI and direct OpenAI are supported transparently via
:func:`create_openai_client`.
"""

from __future__ import annotations

import json
import logging
from typing import Any, Dict, List, Literal, Optional

from pydantic import BaseModel as PydanticBaseModel

from .openai_compat import (
    create_openai_client,
    get_effective_model,
    get_openai_client_kwargs,
    is_azure_openai,
)
from .structured import (
    _resolve_refs,
    build_model_from_schema,
    preprocess_json_schema,
)

logger = logging.getLogger(__name__)


def _enforce_strict_required(schema: Dict[str, Any]) -> Dict[str, Any]:
    """Recursively ensure all object properties are in ``required`` (strict mode)."""
    if not isinstance(schema, dict):
        return schema
    s = dict(schema)
    if s.get("type") == "object" and isinstance(s.get("properties"), dict):
        props = dict(s["properties"])
        s["required"] = list(props.keys())
        for prop_name, prop_schema in props.items():
            props[prop_name] = _enforce_strict_required(prop_schema)
        s["properties"] = props
    if "items" in s and isinstance(s["items"], dict):
        items = _enforce_strict_required(s["items"])
        if "properties" in items and "type" not in items:
            items["type"] = "object"
        elif "items" in items and "type" not in items:
            items["type"] = "array"
        s["items"] = items
    return s


def _build_response_format(
    json_schema: Dict[str, Any],
    schema_name: str,
) -> tuple[Dict[str, Any], Dict[str, Any]]:
    """Preprocess a JSON schema and build the ``response_format`` payload.

    Returns ``(preprocessed_schema, response_format_dict)`` so the caller can
    reuse the preprocessed schema for Pydantic model construction.
    """
    pre = preprocess_json_schema(json_schema, enforce_all_required=True)

    response_format = dict(pre)

    # Resolve $ref if present
    if "$defs" in response_format:
        defs = response_format.pop("$defs")
        response_format = _resolve_refs(response_format, defs)

    if "title" not in response_format:
        response_format["title"] = schema_name

    response_format = _enforce_strict_required(response_format)

    return pre, response_format


async def extract_structured_vision(
    *,
    images: List[str],
    instruction: str,
    json_schema: Dict[str, Any],
    schema_name: str = "VisionSchema",
    system: Optional[str] = None,
    model: Optional[str] = None,
    detail: Literal["auto", "high", "low"] = "auto",
    api_key: Optional[str] = None,
    reasoning_effort: Optional[str] = None,
) -> PydanticBaseModel:
    """Extract structured data from images using a JSON schema.

    Sends one or more images together with a text instruction to the Chat
    Completions API and returns a Pydantic model matching *json_schema*.

    Args:
        images: List of image URLs (``https://...``) or base64 data URIs
            (``data:image/png;base64,...``).  Both formats are passed directly
            to the ``image_url`` content type which accepts either.
        instruction: The text prompt describing what to extract.
        json_schema: JSON schema defining the target structure.
        schema_name: Name identifier for the schema.
        system: Optional system message (a sensible default is used if omitted).
        model: Model override.  Falls back to ``AZURE_OPENAI_DEPLOYMENT_NAME``
            on Azure or ``OPENAI_MODEL`` / ``gpt-4.1-mini`` on direct OpenAI.
        detail: Image detail level — ``"auto"`` (default), ``"high"`` for
            OCR-heavy tasks, or ``"low"`` for cheaper/faster processing.
        api_key: Optional API key override.
        reasoning_effort: Reasoning effort for GPT-5.x models (``"none"``,
            ``"low"``, ``"medium"``, ``"high"``).

    Returns:
        A Pydantic ``BaseModel`` instance whose fields match *json_schema*.
    """
    if not images:
        raise ValueError("At least one image is required")

    # ---- schema preprocessing ------------------------------------------------
    pre, response_format = _build_response_format(json_schema, schema_name)

    logger.info(
        "extract_structured_vision: schema=%s, images=%d, detail=%s, model=%s",
        schema_name,
        len(images),
        detail,
        model,
    )

    # ---- messages -------------------------------------------------------------
    sys_msg = system or (
        "You are a helpful AI assistant that extracts information from images "
        "based on a provided JSON schema. You produce only valid JSON per the "
        "bound schema. If a field is not found in the image, omit it when "
        "optional or set null if allowed. Do not invent values."
    )

    user_content: List[Dict[str, Any]] = [
        {"type": "text", "text": instruction},
    ]
    for img in images:
        user_content.append(
            {"type": "image_url", "image_url": {"url": img, "detail": detail}}
        )

    messages: List[Dict[str, Any]] = [
        {"role": "system", "content": sys_msg},
        {"role": "user", "content": user_content},
    ]

    # ---- API call -------------------------------------------------------------
    effective_model = get_effective_model(model)

    if is_azure_openai():
        client = create_openai_client(api_key=api_key, async_=False)
    else:
        from openai import OpenAI

        client = OpenAI(api_key=api_key, **get_openai_client_kwargs())

    chat_kwargs: Dict[str, Any] = {
        "model": effective_model,
        "messages": messages,
        "response_format": {
            "type": "json_schema",
            "json_schema": {
                "name": schema_name,
                "strict": True,
                "schema": response_format,
            },
        },
    }

    if reasoning_effort and reasoning_effort != "none":
        chat_kwargs["reasoning_effort"] = reasoning_effort

    logger.debug("extract_structured_vision chat_kwargs keys: %s", list(chat_kwargs.keys()))

    resp = client.chat.completions.create(**chat_kwargs)
    output_text = resp.choices[0].message.content

    if not output_text:
        raise ValueError("Vision model returned empty response")

    # ---- parse & return -------------------------------------------------------
    result_dict = json.loads(output_text)

    model_cls = build_model_from_schema(schema_name, pre, base=PydanticBaseModel)
    return model_cls(**result_dict)
