#!/usr/bin/env python3
"""Package model artifacts into an inference cache bundle.

This script creates a lightweight inference bundle containing:
- Model weights (model.pth)
- LLM embeddings (llm_embeddings.pkl)
- LLM descriptions (llm_descriptions.json)
- Config (config.yaml)

Usage:
    python scripts/package_inference_cache.py \
        --checkpoint model_path/sweep_openai_large_512_best.pth \
        --output registry/eb1_router_v1

    # Or use defaults:
    python scripts/package_inference_cache.py
"""

import argparse
import hashlib
import shutil
from datetime import datetime
from pathlib import Path

import yaml


def compute_md5(filepath: Path) -> str:
    """Compute MD5 hash of a file."""
    md5 = hashlib.md5()
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            md5.update(chunk)
    return md5.hexdigest()


def main():
    parser = argparse.ArgumentParser(description="Package inference cache bundle")
    parser.add_argument(
        "--checkpoint",
        type=str,
        default="model_path/hybrid_lambda_rank_openai512_best.pth",
        help="Path to model checkpoint",
    )
    parser.add_argument(
        "--llm-embeddings",
        type=str,
        default="configs/llm_embedding_openai_large_512.pkl",
        help="Path to LLM embeddings pickle",
    )
    parser.add_argument(
        "--llm-descriptions",
        type=str,
        default="configs/LLM_Descriptions.json",
        help="Path to LLM descriptions JSON",
    )
    parser.add_argument(
        "--config",
        type=str,
        default="configs/sweep_openai_large_512.yaml",
        help="Path to training config YAML",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="registry/eb1_router_inference_v1",
        help="Output directory for inference cache",
    )
    args = parser.parse_args()

    # Resolve paths relative to embergraphrouter/
    base_dir = Path(__file__).parent.parent
    checkpoint = base_dir / args.checkpoint
    llm_embeddings = base_dir / args.llm_embeddings
    llm_descriptions = base_dir / args.llm_descriptions
    config_path = base_dir / args.config
    output_dir = base_dir / args.output

    # Validate inputs exist
    for path, name in [
        (checkpoint, "checkpoint"),
        (llm_embeddings, "llm_embeddings"),
        (llm_descriptions, "llm_descriptions"),
    ]:
        if not path.exists():
            raise FileNotFoundError(f"{name} not found: {path}")

    # Create output directory
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"Creating inference cache in: {output_dir}")
    print()

    # Copy model weights
    dst_model = output_dir / "model.pth"
    shutil.copy(checkpoint, dst_model)
    model_size = dst_model.stat().st_size / (1024 * 1024)
    model_md5 = compute_md5(dst_model)
    print(f"  Copied model.pth ({model_size:.1f} MB, md5: {model_md5})")

    # Copy LLM embeddings
    dst_emb = output_dir / "llm_embeddings.pkl"
    shutil.copy(llm_embeddings, dst_emb)
    print(f"  Copied llm_embeddings.pkl")

    # Copy LLM descriptions
    dst_desc = output_dir / "llm_descriptions.json"
    shutil.copy(llm_descriptions, dst_desc)
    print(f"  Copied llm_descriptions.json")

    # Load and save config (subset relevant for inference)
    if config_path.exists():
        with open(config_path, "r") as f:
            full_config = yaml.safe_load(f)

        inference_config = {
            "embed_dim": full_config.get("embed_dim", 512),
            "hidden_dim": full_config.get("hidden_dim", 384),
            "num_gat_layers": full_config.get("num_gat_layers", 3),
            "num_gat_heads": full_config.get("num_gat_heads", 4),
            "dropout": full_config.get("dropout", 0.4),
            "edge_feature_dim": full_config.get("edge_feature_dim", 6),
            "llm_num": full_config.get("llm_num", 7),
            "use_uncertainty_head": full_config.get("use_uncertainty_head", False),
            "use_difficulty_head": full_config.get("use_difficulty_head", False),
        }
    else:
        inference_config = {
            "embed_dim": 512,
            "hidden_dim": 384,
            "num_gat_layers": 3,
            "num_gat_heads": 4,
            "dropout": 0.4,
            "edge_feature_dim": 6,
            "llm_num": 7,
        }

    dst_config = output_dir / "config.yaml"
    with open(dst_config, "w") as f:
        yaml.dump(inference_config, f)
    print(f"  Created config.yaml")

    # Create metadata
    metadata = {
        "name": output_dir.name,
        "version": "1.0.0",
        "created_at": datetime.utcnow().isoformat() + "Z",
        "model_checkpoint": args.checkpoint,
        "model_md5": model_md5,
        "model_size_mb": round(model_size, 2),
        "embedding_source": args.llm_embeddings,
        "embedding_dim": inference_config.get("embed_dim", 512),
        "num_llms": inference_config.get("llm_num", 7),
        "architecture": {
            "type": "HeteroGATRouter",
            "layers": inference_config.get("num_gat_layers", 3),
            "heads": inference_config.get("num_gat_heads", 4),
            "hidden_dim": inference_config.get("hidden_dim", 384),
        },
        "artifacts": [
            "model.pth",
            "llm_embeddings.pkl",
            "llm_descriptions.json",
            "config.yaml",
        ],
        "usage": {
            "load": "RouterInference.from_registry('eb1_router_inference_v1')",
            "route": "router.route(query_embedding)",
        },
    }

    dst_metadata = output_dir / "METADATA.yaml"
    with open(dst_metadata, "w") as f:
        yaml.dump(metadata, f, default_flow_style=False)
    print(f"  Created METADATA.yaml")

    print()
    print("Inference cache created successfully!")
    print()
    print("Usage:")
    print("  from embergraphrouter.inference import RouterInference")
    print(f"  router = RouterInference.from_registry('{output_dir.name}')")
    print("  result = router.route(query_embedding)")
    print("  print(result.best_model, result.confidence)")


if __name__ == "__main__":
    main()
