from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
import urllib.request


def post_json(url: str, payload: dict) -> dict:
    data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=data,
        headers={
            "Content-Type": "application/json",
            "Authorization": "Bearer not-needed",
        },
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=600) as resp:
        return json.loads(resp.read().decode("utf-8"))


def l2_norm(values: list[float]) -> float:
    return math.sqrt(sum(v * v for v in values))


def build_media_part(media_type: str, media_path: str) -> dict:
    if media_type == "image":
        return {"type": "image_url", "image_url": {"url": media_path}}
    if media_type == "audio":
        return {"type": "audio_url", "audio_url": {"url": media_path}}
    if media_type == "video":
        return {"type": "video_url", "video_url": {"url": media_path}}
    raise ValueError(f"Unsupported media_type: {media_type}")


def main() -> None:
    parser = argparse.ArgumentParser(description="OpenAI-compatible multimodal embedding demo for axllm serve")
    parser.add_argument("--model", required=True, type=str)
    parser.add_argument("--api-url", default="http://127.0.0.1:8000/v1", type=str)
    parser.add_argument("--prompt-name", default="query", choices=["query", "document"])
    parser.add_argument("--media-type", required=True, choices=["image", "audio", "video"])
    parser.add_argument("--media-path", required=True, type=Path)
    parser.add_argument("--text", default="", type=str, help="Optional text prefix before the media item")
    args = parser.parse_args()

    content = [{"type": "text", "text": args.text}]
    content.append(build_media_part(args.media_type, str(args.media_path.resolve())))

    payload = {
        "model": args.model,
        "prompt_name": args.prompt_name,
        "encoding_format": "float",
        "messages": [
            {
                "role": "user",
                "content": content,
            }
        ],
    }
    response = post_json(args.api_url.rstrip("/") + "/embeddings", payload)
    emb = response["data"][0]["embedding"]
    head = " ".join(f"{value:.5f}" for value in emb[:8])
    print(f"dim={len(emb)} norm={l2_norm(emb):.6f} head={head}")


if __name__ == "__main__":
    main()