from __future__ import annotations

import argparse
import json
import math
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 main() -> None:
    parser = argparse.ArgumentParser(description="OpenAI-compatible text 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("--input", action="append", default=[], help="Repeatable input text")
    args = parser.parse_args()

    inputs = args.input or [
        "What causes the seasons on Earth?",
        "Earth's axial tilt causes the seasons as the planet orbits the Sun."
    ]
    payload = {
        "model": args.model,
        "input": inputs,
        "prompt_name": args.prompt_name,
        "encoding_format": "float",
    }
    response = post_json(args.api_url.rstrip("/") + "/embeddings", payload)
    print(f"inputs={len(response['data'])}")
    for index, item in enumerate(response["data"]):
        emb = item["embedding"]
        head = " ".join(f"{value:.5f}" for value in emb[:8])
        print(f"[{index}] dim={len(emb)} norm={l2_norm(emb):.6f} head={head}")


if __name__ == "__main__":
    main()