#!/usr/bin/env python3
"""
Sixpert K2 - Vision/Multimodal Example
=======================================
Demonstrates how to use Sixpert K2's vision capabilities with
image inputs. K2's MoE architecture provides specialized experts
for visual understanding.

Usage:
    python vision_example.py --image ./diagram.png --prompt "Analyze this diagram"
"""

import argparse
import base64
import sys

try:
    from llama_cpp import Llama
except ImportError:
    print("Installing llama-cpp-python...")
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "llama-cpp-python"])
    from llama_cpp import Llama


def image_to_base64(image_path: str) -> str:
    """Convert an image file to base64 data URL."""
    with open(image_path, "rb") as f:
        data = base64.b64encode(f.read()).decode("utf-8")

    if image_path.lower().endswith(".png"):
        mime = "image/png"
    elif image_path.lower().endswith(".jpg") or image_path.lower().endswith(".jpeg"):
        mime = "image/jpeg"
    elif image_path.lower().endswith(".webp"):
        mime = "image/webp"
    elif image_path.lower().endswith(".gif"):
        mime = "image/gif"
    else:
        mime = "image/png"

    return f"data:{mime};base64,{data}"


def analyze_image(
    model_path: str,
    image_path: str,
    prompt: str,
    max_tokens: int = 4096,
):
    """Analyze an image using Sixpert K2."""
    print(f"Loading Sixpert K2 (MoE)...")
    llm = Llama(
        model_path=model_path,
        n_ctx=131072,
        n_gpu_layers=-1,
        verbose=False,
    )

    image_data = image_to_base64(image_path)

    messages = [
        {
            "role": "system",
            "content": "You are Sixpert K2, a deep reasoning engine with advanced multimodal capabilities. Analyze images with exceptional detail and precision.",
        },
        {
            "role": "user",
            "content": [
                {"type": "text", "text": prompt},
                {"type": "image_url", "image_url": {"url": image_data}},
            ],
        },
    ]

    print(f"\nPrompt: {prompt}")
    print(f"Image: {image_path}")
    print("-" * 40)
    print("Analyzing with MoE vision experts...")
    print("-" * 40)

    response = llm.create_chat_completion(
        messages=messages,
        max_tokens=max_tokens,
        temperature=0.6,
        stream=True,
    )

    for chunk in response:
        delta = chunk["choices"][0]["delta"].get("content", "")
        if delta:
            print(delta, end="", flush=True)

    print("\n")


def main():
    parser = argparse.ArgumentParser(description="Sixpert K2 Vision Example")
    parser.add_argument("--model", type=str, default="SixpertK2.gguf", help="Path to GGUF model")
    parser.add_argument("--image", type=str, required=True, help="Path to input image")
    parser.add_argument("--prompt", type=str, default="Describe this image in detail.", help="Prompt")
    parser.add_argument("--max-tokens", type=int, default=4096, help="Max tokens")

    args = parser.parse_args()
    analyze_image(args.model, args.image, args.prompt, args.max_tokens)


if __name__ == "__main__":
    main()