import os
import base64
import json
from typing import Any, Union
from openai import OpenAI
from pydantic import BaseModel

# Default base_url pointing to Lemonade's OpenAI-compatible endpoint
API_HOST = os.getenv("OPENAI_API_HOST", "https://lemonade.matthiasg.dev/api/v1")
AI_MODEL = os.getenv("OPENAI_MODEL", "LFM2.5-VL-1.6B-GGUF-Q8_0")
AI_MODEL_THINK = os.getenv("OPENAI_MODEL_THINK", "False").lower() in ("true", "1", "t")
OPENAPI_TOKEN = os.getenv("OPENAI_API_TOKEN", "lemonade")

print(f"Using API Host: {API_HOST}")
print(f"Using API Model: {AI_MODEL}")
print(f"Using API Model Think: {AI_MODEL_THINK}")
print(f"Using API Token: {'Yes' if OPENAPI_TOKEN else 'No'} {OPENAPI_TOKEN[:5] + '...' if OPENAPI_TOKEN else ''}")

# Initialize OpenAI Client for Lemonade
client = OpenAI(
    base_url=API_HOST,
    api_key=OPENAPI_TOKEN if OPENAPI_TOKEN else "lemonade"
)


def generate(
    model: str = AI_MODEL,
    prompt: str = "",
    images: list[str] = None,
    output_format: Union[str, dict, type[BaseModel]] = None,
    system_message: str = None,
    keep_alive: bool = True,
    should_think: bool = AI_MODEL_THINK,
) -> str:
    """
    Generate a completion response using the OpenAI SDK against Lemonade / OpenAI compatible APIs.
    Handles multimodal image content, structured outputs, and custom Lemonade parameters.
    """
    if images is None:
        images = []

    messages = []
    
    # 1. Add System Message if present
    if system_message:
        messages.append({"role": "system", "content": system_message})

    # 2. Build User Content Payload (Text + Multimodal Images)
    user_content = []
    if prompt:
        user_content.append({"type": "text", "text": prompt})

    for img_path in images:
        if os.path.isfile(img_path):
            with open(img_path, "rb") as image_file:
                b64_str = base64.b64encode(image_file.read()).decode("utf-8")
        else:
            b64_str = img_path

        # Ensure base64 string includes Data URI prefix for OpenAI vision format
        if not b64_str.startswith("data:"):
            image_url = f"data:image/jpeg;base64,{b64_str}"
        else:
            image_url = b64_str

        user_content.append({
            "type": "image_url",
            "image_url": {"url": image_url}
        })

    # Simplify content payload if text-only
    if len(user_content) == 1 and user_content[0]["type"] == "text":
        messages.append({"role": "user", "content": prompt})
    else:
        messages.append({"role": "user", "content": user_content})

    # 3. Prepare parameters and custom body flags
    extra_body = {
        "think": should_think,
    }
    if not keep_alive:
        extra_body["keep_alive"] = "0m"

    request_kwargs: dict[str, Any] = {
        "model": model,
        "messages": messages,
        "extra_body": extra_body,
    }

    # 4. Handle Structured Outputs / JSON mode
    if output_format is not None:
        if isinstance(output_format, type) and issubclass(output_format, BaseModel):
            # Pydantic Model -> Use structured output parsing API
            response = client.beta.chat.completions.parse(
                response_format=output_format,
                **request_kwargs
            )
            return response.choices[0].message.content or ""
            
        elif isinstance(output_format, dict):
            # JSON Schema dictionary
            request_kwargs["response_format"] = {
                "type": "json_schema",
                "json_schema": {
                    "name": "structured_response",
                    "strict": True,
                    "schema": output_format,
                }
            }
        elif isinstance(output_format, str) and output_format.lower() == "json":
            # Standard JSON Mode
            request_kwargs["response_format"] = {"type": "json_object"}

    # Standard completion request
    response = client.chat.completions.create(**request_kwargs)
    return response.choices[0].message.content or ""


if __name__ == "__main__":
    result = generate(
        prompt="Explain what you see or answer this prompt.",
        should_think=AI_MODEL_THINK,
    )
    print(result)