#!/usr/bin/env python

"""████
 ██    ██    Datature
   ██  ██    Powering Breakthrough AI
     ██

@File    :   utils.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Utility functions for local deployment server.
"""

import base64
import io
import tempfile
from typing import Any

from PIL import Image
from pydantic import BaseModel, create_model


def json_schema_to_pydantic(
    schema: dict[str, Any], model_name: str = "DynamicModel"
) -> type[BaseModel]:
    """Convert a JSON schema dict to a dynamic Pydantic model.

    This function creates a Pydantic model class at runtime from a JSON schema
    definition, allowing OpenAI-style JSON schemas to be used with Vi's
    response_format parameter.

    Args:
        schema: JSON schema dictionary with properties and types
        model_name: Name for the generated Pydantic model class

    Returns:
        Dynamically created Pydantic BaseModel class

    Example:
        ```python
        schema = {
            "type": "object",
            "properties": {"name": {"type": "string"}, "age": {"type": "integer"}},
            "required": ["name"],
        }

        Model = json_schema_to_pydantic(schema, "Person")
        instance = Model(name="John", age=30)
        ```

    Note:
        This is a simplified conversion that handles common JSON schema patterns.
        Complex schemas with advanced features may require more sophisticated handling.

    """
    if schema.get("type") != "object":
        raise ValueError(
            f"Only 'object' type schemas are supported, got: {schema.get('type')}"
        )

    properties = schema.get("properties", {})
    required = set(schema.get("required", []))

    # Build field definitions for Pydantic
    field_definitions: dict[str, Any] = {}

    for field_name, field_schema in properties.items():
        field_type = _json_type_to_python(field_schema)
        is_required = field_name in required

        # If field is required, just use the type
        # If optional, use type | None with default None
        if is_required:
            field_definitions[field_name] = (field_type, ...)
        else:
            # Make it optional with None default
            if hasattr(field_type, "__origin__"):  # Check if already Optional/Union
                field_definitions[field_name] = (field_type, None)
            else:
                field_definitions[field_name] = (field_type | None, None)

    # Create the Pydantic model dynamically
    return create_model(model_name, **field_definitions)


def _json_type_to_python(field_schema: dict[str, Any]) -> type:
    """Convert JSON schema type to Python type annotation.

    Args:
        field_schema: JSON schema field definition

    Returns:
        Python type or typing annotation

    """
    json_type = field_schema.get("type")

    # Handle array type
    if json_type == "array":
        items_schema = field_schema.get("items", {})
        if items_schema:
            item_type = _json_type_to_python(items_schema)
            return list[item_type]
        return list[Any]

    # Handle object type (nested)
    if json_type == "object":
        # For nested objects, recursively create a Pydantic model
        nested_name = field_schema.get("title", "NestedModel")
        return json_schema_to_pydantic(field_schema, nested_name)

    # Handle enum
    if "enum" in field_schema:
        # For enums, just use str and let xgrammar handle the constraint
        return str

    # Handle primitive types
    type_mapping = {
        "string": str,
        "integer": int,
        "number": float,
        "boolean": bool,
        "null": type(None),
    }

    python_type = type_mapping.get(json_type, Any)

    # Handle number constraints (for validation hints, though xgrammar handles actual constraint)
    if json_type in ("integer", "number"):
        # Pydantic will handle these naturally, xgrammar enforces them
        pass

    return python_type


def convert_response_format_to_pydantic(response_format: dict | None) -> Any | None:
    """Convert OpenAI response_format to Pydantic model.

    Args:
        response_format: OpenAI response format specification, e.g.:
            {
                "type": "json_schema",
                "json_schema": {
                    "name": "response",
                    "schema": {...}
                }
            }

    Returns:
        Pydantic model class, or None if not applicable

    """
    if not response_format or response_format.get("type") != "json_schema":
        return None

    # Extract schema from OpenAI format
    json_schema_obj = response_format.get("json_schema", {})
    schema = json_schema_obj.get("schema", {})
    model_name = json_schema_obj.get("name", "CustomResponse")

    if not schema:
        return None

    # Convert JSON schema to Pydantic model
    return json_schema_to_pydantic(schema, model_name)


def convert_messages_to_openai_format(
    messages: list[dict] | None, encode_images: bool = False
) -> list[dict]:
    """Convert Vi SDK messages to OpenAI format.

    This is a shared utility function used by multiple backends to convert
    messages from Vi SDK format to OpenAI chat format.

    Args:
        messages: List of chat messages in Vi SDK format
        encode_images: Whether to encode local image paths to base64 data URIs
                      (for backends that support it like SGLang)

    Returns:
        List of messages in OpenAI chat format

    """
    if not messages:
        return []

    openai_messages = []
    for message in messages:
        role = message.get("role", "user")
        content = message.get("content", "")

        # Handle different content formats
        if isinstance(content, str):
            openai_messages.append({"role": role, "content": content})
        elif isinstance(content, list):
            # Handle multimodal content (text + images)
            openai_content = []
            for item in content:
                if isinstance(item, dict):
                    if item.get("type") == "text":
                        openai_content.append(
                            {"type": "text", "text": item.get("text", "")}
                        )
                    elif item.get("type") == "image":
                        # Vi SDK extension - local image path
                        image_path = item.get("image")
                        if image_path:
                            if encode_images:
                                # Convert local path to base64 data URI
                                try:
                                    image_data = encode_image_to_base64(image_path)
                                    openai_content.append(
                                        {
                                            "type": "image_url",
                                            "image_url": {
                                                "url": f"data:image/jpeg;base64,{image_data}"
                                            },
                                        }
                                    )
                                except Exception:
                                    # Fallback to file:// URL if encoding fails
                                    openai_content.append(
                                        {
                                            "type": "image_url",
                                            "image_url": {
                                                "url": f"file://{image_path}"
                                            },
                                        }
                                    )
                            else:
                                # Use file:// URL directly (for VLLM)
                                openai_content.append(
                                    {
                                        "type": "image_url",
                                        "image_url": {"url": f"file://{image_path}"},
                                    }
                                )
                    elif item.get("type") == "image_url":
                        # OpenAI format - image URL
                        openai_content.append(item)
            if openai_content:
                openai_messages.append({"role": role, "content": openai_content})
        else:
            # Fallback for other content types
            openai_messages.append({"role": role, "content": str(content)})

    return openai_messages


def extract_text_from_messages(messages: list[dict] | None) -> str:
    """Extract text content from OpenAI-formatted messages.

    This is a shared utility function used by multiple backends to extract
    text content from messages in OpenAI format.

    Args:
        messages: List of chat messages in OpenAI format

    Returns:
        Combined text content from all user messages

    """
    if not messages:
        return ""

    text_parts = []
    for message in messages:
        content = message.get("content", "")
        role = message.get("role", "")

        if role == "user":
            if isinstance(content, str):
                text_parts.append(content)
            elif isinstance(content, list):
                # Extract text from multimodal content
                for item in content:
                    if isinstance(item, dict) and item.get("type") == "text":
                        text_parts.append(item.get("text", ""))

    return " ".join(text_parts).strip()


def encode_image_to_base64(image_path: str) -> str:
    """Encode an image file to base64.

    This is a shared utility function for encoding images to base64,
    used by backends that need to send images as data URIs.

    Args:
        image_path: Path to the image file

    Returns:
        Base64 encoded image data

    Raises:
        FileNotFoundError: If image file doesn't exist
        ValueError: If image format is unsupported

    """
    try:
        # Open and convert image to RGB if necessary
        with Image.open(image_path) as img:
            if img.mode != "RGB":
                img = img.convert("RGB")

            # Save to bytes buffer
            buffer = io.BytesIO()
            img.save(buffer, format="JPEG")
            buffer.seek(0)

            # Encode to base64
            return base64.b64encode(buffer.getvalue()).decode("utf-8")
    except ImportError:
        # PIL not available, try direct file reading
        with open(image_path, "rb") as f:
            return base64.b64encode(f.read()).decode("utf-8")


def save_base64_image_to_temp_file(data_uri: str) -> str:
    """Save a base64-encoded image to a temporary file.

    This is a shared utility function for saving base64 images to temp files,
    used by backends that need to work with local file paths.

    Args:
        data_uri: Data URI with base64 image (e.g., "data:image/jpeg;base64,...")

    Returns:
        Path to the temporary file

    Raises:
        ValueError: If the data URI format is invalid

    """
    # Extract base64 data
    if "," in data_uri:
        base64_data = data_uri.split(",", 1)[1]
    else:
        base64_data = data_uri

    # Decode and save to temp file
    image_data = base64.b64decode(base64_data)
    image = Image.open(io.BytesIO(image_data))

    # Create temp file
    temp_file = tempfile.NamedTemporaryFile(
        delete=False, suffix=f".{image.format.lower() if image.format else 'png'}"
    )
    image.save(temp_file.name)
    temp_file.close()

    return temp_file.name


def extract_image_and_prompt_from_messages(
    messages: list[dict] | None,
) -> tuple[str | None, str | None]:
    """Extract image path and text prompt from OpenAI-formatted messages.

    This is a shared utility function for extracting image and text content
    from messages in OpenAI format, used by backends that need to convert
    OpenAI messages to Vi SDK format.

    Args:
        messages: List of chat messages in OpenAI format

    Returns:
        Tuple of (image_path, text_prompt). Both can be None.

    """
    if not messages:
        return None, None

    image_path = None
    prompt_parts = []

    for message in messages:
        content = message.get("content", "")

        if isinstance(content, str):
            # Simple text message
            if message["role"] == "system":
                prompt_parts.append(f"System: {content}")
            elif message["role"] == "user":
                prompt_parts.append(content)
        elif isinstance(content, list):
            # Multimodal message with images and text
            for item in content:
                if isinstance(item, dict):
                    if item.get("type") == "text":
                        prompt_parts.append(item.get("text", ""))
                    elif item.get("type") == "image":
                        # Local image path (Vi SDK extension)
                        image_path = item.get("image")
                    elif item.get("type") == "image_url":
                        # Image URL or base64 (OpenAI format)
                        image_url_obj = item.get("image_url", {})
                        if isinstance(image_url_obj, dict):
                            image_url = image_url_obj.get("url", "")
                        else:
                            # Handle case where image_url is a string directly
                            image_url = str(image_url_obj)

                        if image_url:
                            if image_url.startswith("data:image"):
                                # Base64 encoded image
                                image_path = save_base64_image_to_temp_file(image_url)
                            else:
                                # Regular URL or file path
                                image_path = image_url

    prompt = " ".join(prompt_parts).strip() if prompt_parts else None
    return image_path, prompt


def validate_vi_format_request(source: str | None, user_prompt: str | None) -> None:
    """Validate that a Vi format request has required parameters.

    Args:
        source: Image source path (can be None for text-only requests)
        user_prompt: Text prompt (can be None for image-only requests)

    Raises:
        ValueError: If neither source nor user_prompt is provided

    """
    if not source and not user_prompt:
        raise ValueError(
            "Either 'source' (image path) or 'user_prompt' (text) must be provided"
        )


def build_generation_config_from_request(request) -> dict[str, Any]:
    """Build generation config dictionary from a ChatCompletionRequest.

    This function extracts generation parameters from a ChatCompletionRequest
    and builds a generation config dictionary in the format expected by backends.

    Args:
        request: ChatCompletionRequest object with generation parameters

    Returns:
        Dictionary with generation configuration parameters

    """
    generation_config = {}

    # Start with any user-provided generation_config
    if hasattr(request, "generation_config") and request.generation_config:
        generation_config.update(request.generation_config)

    # Override with individual parameters (they take precedence)
    temperature = getattr(request, "temperature", None)
    if temperature is not None:
        generation_config["temperature"] = temperature
        generation_config["do_sample"] = temperature > 0.0
    else:
        # Default temperature
        generation_config["temperature"] = 0.7
        generation_config["do_sample"] = True

    # Map other parameters
    parameter_mappings = {
        "max_tokens": "max_new_tokens",
        "top_p": "top_p",
        "top_k": "top_k",
        "seed": "seed",
        "repetition_penalty": "repetition_penalty",
        "stop": "stop",
        "logit_bias": "logit_bias",
    }

    for request_attr, config_key in parameter_mappings.items():
        value = getattr(request, request_attr, None)
        if value is not None:
            generation_config[config_key] = value

    # Handle OpenAI-style penalties (convert to repetition_penalty if needed)
    frequency_penalty = getattr(request, "frequency_penalty", None)
    presence_penalty = getattr(request, "presence_penalty", None)

    if frequency_penalty is not None or presence_penalty is not None:
        if "repetition_penalty" not in generation_config:
            combined_penalty = 1.0
            if frequency_penalty is not None:
                combined_penalty += frequency_penalty * 0.1
            if presence_penalty is not None:
                combined_penalty += presence_penalty * 0.1
            if combined_penalty != 1.0:
                generation_config["repetition_penalty"] = combined_penalty

    return generation_config


def handle_vi_backend_source_error(
    error: ValueError, source: str | None, user_prompt: str | None
) -> ValueError:
    """Handle source parameter validation errors for Vi-compatible backends.

    This function provides more helpful error messages when a model requires
    an image source but none was provided.

    Args:
        error: The original ValueError
        source: Image source path that was provided
        user_prompt: Text prompt that was provided

    Returns:
        ValueError with enhanced error message

    """
    error_msg = str(error)
    if "source parameter is required" in error_msg:
        if source is None:
            error_message = (
                "This model requires an image source, but none was provided.\n"
                f"Extracted from request: source={source}, user_prompt={repr(user_prompt)}\n\n"
                "To provide an image:\n"
                "  - OpenAI format: Use multimodal content array:\n"
                "    {\n"
                '      "messages": [{\n'
                '        "role": "user",\n'
                '        "content": [\n'
                '          {"type": "image_url", "image_url": {"url": "/path/to/image.jpg"}},\n'
                '          {"type": "text", "text": "Your question"}\n'
                "        ]\n"
                "      }]\n"
                "    }\n\n"
                "  - Vi SDK format: Use source field directly:\n"
                '    {"source": "/path/to/image.jpg", "user_prompt": "Your question"}'
            )
            return ValueError(error_message)
    return error
