"""
Dokodemo AI CLI

Usage:
    dokodemo run <model> --prompt "Hello, world!" [options]
    dokodemo benchmark <model> [options]
    dokodemo profile <model> --prompt "..." [options]
    dokodemo info <model>
"""

from __future__ import annotations

import argparse
import logging
import sys
import time

logger = logging.getLogger(__name__)


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="dokodemo",
        description=(
            "Dokodemo AI - Run trillion-parameter models on consumer GPUs\n"
            "  どこでも AI - どんなGPUでも大規模モデルを動かす"
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    sub = parser.add_subparsers(dest="command", required=True)

    # ---- Common arguments ----
    def add_common(p: argparse.ArgumentParser) -> None:
        p.add_argument("model", help="HuggingFace model ID or local path")
        p.add_argument(
            "--compression",
            choices=["4bit", "8bit", "2bit", "dynamic"],
            default=None,
            help="Weight compression (default: none / FP16)",
        )
        p.add_argument(
            "--max-gpu-memory",
            default=None,
            metavar="SIZE",
            help="Max GPU memory for weights, e.g. '4GB', '2048MB'",
        )
        p.add_argument(
            "--device",
            default=None,
            help="Device: 'cuda', 'cuda:1', 'mps', 'cpu'",
        )
        p.add_argument(
            "--token",
            default=None,
            metavar="HF_TOKEN",
            help="HuggingFace API token for gated models",
        )
        p.add_argument(
            "--trust-remote-code",
            action="store_true",
            help="Allow running remote model code",
        )

    # ---- run subcommand ----
    run_p = sub.add_parser("run", help="Generate text from a prompt")
    add_common(run_p)
    run_p.add_argument(
        "--prompt", "-p", required=True, help="Input prompt text"
    )
    run_p.add_argument(
        "--max-new-tokens", "-n", type=int, default=100,
        help="Maximum tokens to generate (default: 100)",
    )
    run_p.add_argument(
        "--temperature", type=float, default=0.7,
        help="Sampling temperature (default: 0.7)",
    )
    run_p.add_argument(
        "--top-k", type=int, default=50,
        help="Top-k sampling (default: 50)",
    )
    run_p.add_argument(
        "--top-p", type=float, default=0.9,
        help="Nucleus sampling p (default: 0.9)",
    )
    run_p.add_argument(
        "--greedy", action="store_true",
        help="Use greedy decoding instead of sampling",
    )
    run_p.add_argument(
        "--profiling", action="store_true",
        help="Enable detailed performance profiling",
    )

    # ---- benchmark subcommand ----
    bench_p = sub.add_parser(
        "benchmark", help="Measure inference speed"
    )
    add_common(bench_p)
    bench_p.add_argument(
        "--tokens", type=int, default=50,
        help="Number of tokens to generate for timing (default: 50)",
    )
    bench_p.add_argument(
        "--warmup", type=int, default=5,
        help="Warmup tokens before timing (default: 5)",
    )

    # ---- profile subcommand ----
    prof_p = sub.add_parser(
        "profile", help="Profile layer-by-layer I/O and compute"
    )
    add_common(prof_p)
    prof_p.add_argument(
        "--prompt", "-p", required=True, help="Input prompt"
    )
    prof_p.add_argument(
        "--tokens", type=int, default=10,
        help="Tokens to generate for profiling (default: 10)",
    )

    # ---- info subcommand ----
    info_p = sub.add_parser(
        "info", help="Show model structure without running inference"
    )
    info_p.add_argument("model", help="HuggingFace model ID or local path")
    info_p.add_argument("--token", default=None)
    info_p.add_argument("--trust-remote-code", action="store_true")

    return parser


def _load_model(args) -> tuple:
    """Load model and tokenizer from parsed args."""
    from dokodemo_ai import AutoModel

    model = AutoModel.from_pretrained(
        args.model,
        compression=getattr(args, "compression", None),
        max_gpu_memory=getattr(args, "max_gpu_memory", None),
        device=getattr(args, "device", None),
        hf_token=getattr(args, "token", None),
        trust_remote_code=getattr(args, "trust_remote_code", False),
        profiling=getattr(args, "profiling", False),
    )
    return model


def cmd_run(args) -> int:
    """Execute the 'run' command."""
    import torch

    model = _load_model(args)

    # Tokenize
    inputs = model.tokenizer(args.prompt, return_tensors="pt")
    input_ids = inputs["input_ids"]

    print(f"\n{'='*60}")
    print(f"Prompt: {args.prompt}")
    print(f"{'='*60}")
    print("Generating", end="", flush=True)

    t0 = time.perf_counter()
    with torch.no_grad():
        output = model.generate(
            input_ids=input_ids,
            max_new_tokens=args.max_new_tokens,
            temperature=args.temperature,
            top_k=args.top_k,
            top_p=args.top_p,
            do_sample=not args.greedy,
        )
    elapsed = time.perf_counter() - t0

    new_tokens = output.shape[1] - input_ids.shape[1]
    tps = new_tokens / elapsed

    # Decode
    full_text = model.tokenizer.decode(output[0], skip_special_tokens=True)
    generated_text = full_text[len(args.prompt):]

    print(f"\r{'='*60}")
    print(generated_text)
    print(f"{'='*60}")
    print(
        f"Generated {new_tokens} tokens in {elapsed:.1f}s "
        f"({tps:.2f} tokens/sec)"
    )

    if getattr(args, "profiling", False):
        print("\n" + model.get_profiling_report())

    return 0


def cmd_benchmark(args) -> int:
    """Execute the 'benchmark' command."""
    import torch

    model = _load_model(args)

    # Warmup with a short prompt
    warmup_prompt = "The quick brown fox"
    inputs = model.tokenizer(warmup_prompt, return_tensors="pt")
    input_ids = inputs["input_ids"]

    print(f"Warming up ({args.warmup} tokens)...", end="", flush=True)
    with torch.no_grad():
        model.generate(input_ids=input_ids, max_new_tokens=args.warmup)
    print(" done")

    # Benchmark
    bench_prompt = "In the field of artificial intelligence,"
    inputs = model.tokenizer(bench_prompt, return_tensors="pt")
    input_ids = inputs["input_ids"]

    print(f"Benchmarking ({args.tokens} tokens)...", end="", flush=True)
    t0 = time.perf_counter()
    with torch.no_grad():
        output = model.generate(
            input_ids=input_ids,
            max_new_tokens=args.tokens,
            do_sample=False,
        )
    elapsed = time.perf_counter() - t0
    print(" done")

    new_tokens = output.shape[1] - input_ids.shape[1]
    tps = new_tokens / elapsed
    ms_per_token = 1000 * elapsed / new_tokens

    print(f"\n{'='*60}")
    print("BENCHMARK RESULTS")
    print(f"{'='*60}")
    print(f"Model:           {args.model}")
    print(f"Compression:     {args.compression or 'none (FP16)'}")
    print(f"Tokens generated:{new_tokens}")
    print(f"Total time:      {elapsed:.2f}s")
    print(f"Throughput:      {tps:.3f} tokens/sec")
    print(f"Latency:         {ms_per_token:.0f} ms/token")

    if model._graph.is_moe:
        stats = model.get_expert_stats()
        if stats:
            print(f"\nMoE Expert Stats:")
            for layer_key, data in list(stats.items())[:3]:
                print(
                    f"  {layer_key}: {data['unique_experts_used']} unique "
                    f"experts used out of {model._graph.layers[0].num_experts}"
                )

    return 0


def cmd_profile(args) -> int:
    """Execute the 'profile' command."""
    import torch

    # Force profiling on
    args.profiling = True
    model = _load_model(args)

    inputs = model.tokenizer(args.prompt, return_tensors="pt")
    input_ids = inputs["input_ids"]

    with torch.no_grad():
        model.generate(
            input_ids=input_ids,
            max_new_tokens=args.tokens,
        )

    print(model.get_profiling_report())
    return 0


def cmd_info(args) -> int:
    """Execute the 'info' command — show model structure."""
    from dokodemo_ai.graph.compiler import ModelGraphCompiler

    print(f"Analyzing model: {args.model}")
    compiler = ModelGraphCompiler()
    graph = compiler.compile(
        args.model,
        hf_token=getattr(args, "token", None),
        trust_remote_code=getattr(args, "trust_remote_code", False),
    )
    print(f"\n{'='*60}")
    print("MODEL STRUCTURE")
    print(f"{'='*60}")
    print(graph.summary())
    print(f"\nLayers ({graph.num_layers} total):")
    for i, layer in enumerate(graph.layers[:5]):
        print(f"  Layer {i}: {layer.size_bytes / 1e6:.1f} MB", end="")
        if layer.is_moe:
            print(
                f"  [MoE: {layer.num_experts} experts, "
                f"{layer.num_active_experts} active per token]"
            )
        else:
            print()
    if graph.num_layers > 5:
        print(f"  ... ({graph.num_layers - 5} more layers)")
    print(f"  (all similar in size)")
    return 0


_COMMANDS = {
    "run": cmd_run,
    "benchmark": cmd_benchmark,
    "profile": cmd_profile,
    "info": cmd_info,
}


def main() -> None:
    parser = _build_parser()
    args = parser.parse_args()

    try:
        rc = _COMMANDS[args.command](args)
        sys.exit(rc)
    except KeyboardInterrupt:
        print("\nInterrupted.")
        sys.exit(1)
    except Exception as e:
        logger.error("Error: %s", e, exc_info=True)
        sys.exit(2)


if __name__ == "__main__":
    main()
