from __future__ import annotations

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

import numpy as np

PACKAGE_ROOT = Path(__file__).resolve().parent.parent
CASE_ROOT = PACKAGE_ROOT / "python" / "testdata" / "service_cases"
DEFAULT_API_PACKAGE_ROOT = PACKAGE_ROOT
DEFAULT_OUTPUT_DIR = PACKAGE_ROOT / "python" / "testdata" / "openai_api_hf_multimodal_compare"
DEFAULT_CASES = (
    ("T0", "embedding_doc"),
    ("T1", "red_planet_query"),
    ("A", "vision_sample"),
    ("B8", "audio_test_chunk0_8s_wav"),
    ("B30", "audio_test_chunk0_30s_wav"),
    ("C", "video_visual_red_panda_openai_mp4"),
)
CASE_LABEL_NAMES = {
    "T0": "text_document",
    "T1": "text_query",
    "A": "image",
    "B8": "audio_8s",
    "B30": "audio_30s",
    "C": "video",
}


def log(message: str) -> None:
    print(f"[compare] {message}", flush=True)


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


def get_default_model(api_url: str) -> str:
    with urllib.request.urlopen(api_url.rstrip("/") + "/models", timeout=60) as response:
        payload = json.loads(response.read().decode("utf-8"))
    data = payload.get("data") or []
    if not data:
        raise RuntimeError("No model found from /v1/models")
    return str(data[0]["id"])


def cosine_similarity(lhs: np.ndarray, rhs: np.ndarray) -> float:
    lhs64 = lhs.reshape(-1).astype(np.float64)
    rhs64 = rhs.reshape(-1).astype(np.float64)
    denom = (np.linalg.norm(lhs64) * np.linalg.norm(rhs64)) + 1e-12
    return float(np.dot(lhs64, rhs64) / denom)


def compare_embeddings(reference: np.ndarray, output: np.ndarray) -> dict:
    diff = np.abs(reference - output)
    return {
        "reference_shape": list(reference.shape),
        "output_shape": list(output.shape),
        "max_abs_diff": float(diff.max()),
        "mean_abs_diff": float(diff.mean()),
        "cosine_similarity": cosine_similarity(reference, output),
    }


def format_tensor_preview(tensor: np.ndarray, limit: int = 8) -> str:
    flat = tensor.reshape(-1).astype(np.float32)
    shown = ", ".join(f"{value:.6f}" for value in flat[:limit])
    suffix = ", ..." if flat.size > limit else ""
    return f"[{shown}{suffix}]"


def build_api_request(model: str, case_meta: dict, api_package_root: Path) -> dict:
    modality = case_meta["modality"]
    prompt_name = case_meta["prompt_name"]

    if modality == "text":
        return {
            "model": model,
            "input": case_meta["text"],
            "prompt_name": prompt_name,
            "encoding_format": "float",
        }

    asset_path = api_package_root / case_meta["asset_path"]

    if modality == "image":
        media_part = {"type": "image_url", "image_url": {"url": str(asset_path)}}
    elif modality == "audio":
        media_part = {"type": "audio_url", "audio_url": {"url": str(asset_path)}}
    elif modality == "video":
        media_part = {"type": "video_url", "video_url": {"url": str(asset_path)}}
    else:
        raise ValueError(f"Unsupported modality for this script: {modality}")

    return {
        "model": model,
        "prompt_name": prompt_name,
        "encoding_format": "float",
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": case_meta.get("text_prefix", "")},
                    media_part,
                ],
            }
        ],
    }


def save_api_only_output(output_dir: Path, label: str, modality: str, api_embedding: np.ndarray) -> dict:
    api_path = output_dir / f"{label}_{modality}_openai_api.npy"
    np.save(api_path, api_embedding.astype(np.float32))
    return {
        "openai_api_npy": str(api_path),
    }


def load_reference_embedding(case_dir: Path) -> np.ndarray:
    hf_path = case_dir / "torch_embedding.npy"
    if not hf_path.exists():
        raise FileNotFoundError(
            f"Missing cached HF reference: {hf_path}. "
            "Generate it once on the server with the original HF model, then package it with the same service-case input. "
            "Board-side comparison must not depend on original HF weights."
        )
    return np.load(hf_path).astype(np.float32).reshape(1, -1)


def print_case_summary(
    *,
    index: int,
    total: int,
    label: str,
    display_name: str,
    case_name: str,
    modality: str,
    prompt_name: str,
    input_desc: str,
    soft_token_count: int | None,
    comparison: dict,
    output_preview: str,
    reference_preview: str,
) -> None:
    log(
        f"[{index}/{total}] {display_name} ({label}) modality={modality} "
        f"prompt={prompt_name} case={case_name}"
    )
    if soft_token_count is not None:
        log(f"  soft_tokens={soft_token_count}")
    log(f"  input={input_desc}")
    log(
        "  output_shape={} cosine={:.6f} mean_abs_diff={:.6f} max_abs_diff={:.6f}".format(
            comparison["output_shape"],
            comparison["cosine_similarity"],
            comparison["mean_abs_diff"],
            comparison["max_abs_diff"],
        )
    )
    log(f"  api_tensor={output_preview}")
    log(f"  hf_tensor={reference_preview}")


def print_final_summary(
    *,
    model_name: str,
    api_url: str,
    summary_path: Path,
    api_npz_path: Path,
    hf_npz_path: Path,
    summary_cases: list[dict],
) -> None:
    log("done")
    log(f"  api_url={api_url}")
    log(f"  model={model_name}")
    for item in summary_cases:
        comparison = item["comparison"]
        log(
            "  {} ({}) cosine={:.6f} mean_abs_diff={:.6f} max_abs_diff={:.6f}".format(
                item["display_name"],
                item["label"],
                comparison["cosine_similarity"],
                comparison["mean_abs_diff"],
                comparison["max_abs_diff"],
            )
        )
    log(f"  summary_json={summary_path}")
    log(f"  api_npz={api_npz_path}")
    log(f"  hf_reference_npz={hf_npz_path}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Run text and multimodal embedding comparison: board-side OpenAI-compatible API vs packaged HF reference embeddings"
    )
    parser.add_argument("--api-url", default="http://127.0.0.1:8000/v1", type=str)
    parser.add_argument("--model", default=None, type=str)
    parser.add_argument(
        "--api-package-root",
        type=Path,
        default=DEFAULT_API_PACKAGE_ROOT,
        help="Board-visible package root used to build media paths inside OpenAI API requests.",
    )
    parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_DIR)
    args = parser.parse_args()

    api_package_root = args.api_package_root
    output_dir = args.output_dir.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    log(f"api_url={args.api_url}")
    log(f"api_package_root={api_package_root}")
    log(f"reference_dir={CASE_ROOT}")
    log("hf_reference_source=packaged torch_embedding.npy files generated once on server")
    log("board_runtime_dependency=axllm service + packaged assets only; original HF safetensors are not used")

    model_name = args.model or get_default_model(args.api_url)
    log(f"model={model_name}")
    api_npz = {}
    hf_npz = {}
    summary_cases = []
    total_cases = len(DEFAULT_CASES)

    for index, (label, case_name) in enumerate(DEFAULT_CASES, start=1):
        case_dir = CASE_ROOT / case_name
        case_meta = json.loads((case_dir / "meta.json").read_text(encoding="utf-8"))
        display_name = CASE_LABEL_NAMES.get(label, label)
        if case_meta["modality"] == "text":
            text = str(case_meta["text"])
            input_desc = text
        else:
            input_desc = str(api_package_root / case_meta["asset_path"])
        payload = build_api_request(model_name, case_meta, api_package_root=api_package_root)
        response = post_json(args.api_url.rstrip("/") + "/embeddings", payload)
        api_embedding = np.asarray(response["data"][0]["embedding"], dtype=np.float32).reshape(1, -1)
        hf_embedding = load_reference_embedding(case_dir)
        comparison = compare_embeddings(hf_embedding, api_embedding)
        output_preview = format_tensor_preview(api_embedding)
        reference_preview = format_tensor_preview(hf_embedding)

        api_npz[label] = api_embedding.astype(np.float32)
        hf_npz[label] = hf_embedding.astype(np.float32)
        file_paths = save_api_only_output(output_dir, label, case_meta["modality"], api_embedding)
        case_summary = {
            "label": label,
            "display_name": display_name,
            "case_name": case_name,
            "modality": case_meta["modality"],
            "prompt_name": case_meta["prompt_name"],
            "comparison": comparison,
            "output_preview": output_preview,
            "reference_preview": reference_preview,
            "hf_meta": {
                "source": "packaged_cached_reference",
                "file": str(case_dir / "torch_embedding.npy"),
                "generation": "server-side HF model inference, pre-generated before board validation",
            },
            **file_paths,
        }
        summary_cases.append(case_summary)
        print_case_summary(
            index=index,
            total=total_cases,
            label=label,
            display_name=display_name,
            case_name=case_name,
            modality=case_meta["modality"],
            prompt_name=case_meta["prompt_name"],
            input_desc=input_desc,
            soft_token_count=case_meta.get("soft_token_count"),
            comparison=comparison,
            output_preview=output_preview,
            reference_preview=reference_preview,
        )

    api_npz_path = output_dir / "openai_api_tensors_cases.npz"
    hf_npz_path = output_dir / "hf_torch_tensors_cases.npz"
    np.savez(api_npz_path, **api_npz)
    np.savez(hf_npz_path, **hf_npz)

    summary = {
        "api_url": args.api_url,
        "model": model_name,
        "api_package_root": str(api_package_root),
        "reference_dir": str(CASE_ROOT),
        "labels": CASE_LABEL_NAMES,
        "hf_reference_source": "packaged_cached_reference",
        "hf_runtime_on_board": False,
        "openai_api_npz": str(api_npz_path),
        "hf_torch_npz": str(hf_npz_path),
        "cases": summary_cases,
    }
    summary_path = output_dir / "summary.json"
    summary_path.write_text(json.dumps(summary, indent=2, ensure_ascii=False), encoding="utf-8")
    print_final_summary(
        model_name=model_name,
        api_url=args.api_url,
        summary_path=summary_path,
        api_npz_path=api_npz_path,
        hf_npz_path=hf_npz_path,
        summary_cases=summary_cases,
    )


if __name__ == "__main__":
    main()