"""Generate placeholder model artifacts with random weights.

This script initializes the router MLP with random weights and saves all
necessary artifacts for the inference harness.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path


def generate_placeholder_artifacts(output_dir: Path, seed: int = 42) -> None:
    """Generate placeholder model artifacts.

    Args:
        output_dir: Directory to write artifacts to.
        seed: Random seed for reproducibility.
    """
    import torch
    import yaml

    from partner_inference_harness.model.router_mlp import (
        RouterConfig,
        RouterMLP,
        save_model,
    )

    # Set seed for reproducibility
    torch.manual_seed(seed)

    # Default model list (generic placeholders)
    model_list = [
        "model_a",
        "model_b",
        "model_c",
        "model_d",
        "model_e",
        "model_f",
        "model_g",
    ]

    # Create config
    config = RouterConfig(
        input_dim=256,
        hidden_dim=2048,
        num_layers=4,
        num_models=len(model_list),
        dropout=0.1,
    )

    # Initialize model with random weights
    model = RouterMLP(config)

    print(f"Model initialized with {model.count_parameters():,} parameters")
    print(f"Memory footprint: {model.get_memory_footprint_mb():.2f} MB")

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

    # Save model weights
    model_path = output_dir / "model.pth"
    save_model(model, model_path)
    print(f"Saved model to {model_path}")

    # Save config
    config_path = output_dir / "config.yaml"
    with open(config_path, "w") as f:
        yaml.dump(
            {
                "version": "1.0",
                "model": {
                    "input_dim": config.input_dim,
                    "hidden_dim": config.hidden_dim,
                    "num_layers": config.num_layers,
                    "num_models": config.num_models,
                },
                "inference": {
                    "batch_size": 32,
                    "max_sequence_length": 512,
                },
            },
            f,
            default_flow_style=False,
        )
    print(f"Saved config to {config_path}")

    # Save model list
    llm_list_path = output_dir / "llm_list.json"
    with open(llm_list_path, "w") as f:
        json.dump({"models": model_list}, f, indent=2)
    print(f"Saved model list to {llm_list_path}")

    print("\nArtifacts generated successfully!")
    print(f"Output directory: {output_dir}")


def main() -> int:
    """CLI entry point."""
    parser = argparse.ArgumentParser(
        description="Generate placeholder model artifacts"
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path(__file__).parent.parent / "artifacts" / "placeholder",
        help="Output directory for artifacts",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for weight initialization",
    )

    args = parser.parse_args()

    try:
        generate_placeholder_artifacts(args.output_dir, args.seed)
        return 0
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())
