import argparse


def main():
    p = argparse.ArgumentParser(description="DiffuThink v2 — from scratch")
    sub = p.add_subparsers(dest="command", required=True)
    prep = sub.add_parser("prepare")
    prep.add_argument("--output", default="data/processed/stories-v2")
    prep.add_argument("--revision", default="f54c09fd23315a6f9c86f9dc80f725de7d8f9c64")
    prep.add_argument("--tokenizer", help="Reuse our existing BPE vocabulary for weight continuity")
    for key, value in [("stories",100000),("heldout",1000),("tokenizer-stories",50000),("vocab-size",8192),("length",192)]:
        prep.add_argument("--"+key, type=int, default=value)
    train = sub.add_parser("train")
    train.add_argument("--data", default="data/processed/stories-v2")
    train.add_argument("--output", default="runs/stories-v2")
    train.add_argument("--resume")
    train.add_argument("--initialize-from",help="Start a new optimization phase from our own trained weights")
    train.add_argument("--device", choices=["cpu","cuda"], default="cuda")
    train.add_argument("--precision", choices=["fp32","bf16"], default="bf16")
    train.add_argument("--lr", type=float, default=0.0006)
    for key, value in [("steps",6000),("batch-size",16),("accumulation",2),("width",320),("heads",8),("layers",6),
                       ("eval-every",500),("eval-samples",128),("test-samples",512),("threads",4),("seed",42),("warmup",300)]:
        train.add_argument("--"+key,type=int,default=value)
    sample = sub.add_parser("sample")
    sample.add_argument("--model", default="runs/stories-512/best")
    sample.add_argument("--generation-mode", choices=["auto", "causal", "diffusion"], default="auto")
    sample.add_argument("--prompt", default="Once upon a time, a little girl found")
    sample.add_argument("--suffix", default=None)
    sample.add_argument("--missing-tokens", type=int, default=3)
    sample.add_argument("--new-tokens", type=int, default=64)
    sample.add_argument("--finish-sentence-tokens",type=int,default=0,help="Optional extra causal tokens (0..64) to reach terminal punctuation")
    sample.add_argument("--block-size", type=int, default=16)
    sample.add_argument("--steps", type=int, default=16)
    sample.add_argument("--temperature", type=float, default=0.5)
    sample.add_argument("--seed", type=int, default=42)
    sample.add_argument("--device", choices=["cpu","cuda"], default="cpu")
    benchmark = sub.add_parser("benchmark")
    benchmark.add_argument("--model",default="runs/stories-v2-refined/best")
    benchmark.add_argument("--data",default="data/processed/stories-v2")
    benchmark.add_argument("--output",default="reports/v2/sampling_validation.json")
    benchmark.add_argument("--device",choices=["cpu","cuda"],default="cuda")
    benchmark.add_argument("--precision",choices=["fp32","bf16"],default="bf16")
    benchmark.add_argument("--split",choices=["validation","test"],default="validation")
    benchmark.add_argument("--examples",type=int,default=64)
    serve = sub.add_parser("serve")
    serve.add_argument("--model",default="runs/stories-512/best")
    serve.add_argument("--device",choices=["cpu","cuda"],default="cuda")
    serve.add_argument("--port",type=int,default=7861)
    export = sub.add_parser("export")
    export.add_argument("--model",default="runs/stories-512/best")
    export.add_argument("--report",default="runs/stories-512/report.json")
    export.add_argument("--output",default="artifacts/DiffuThink-Story-512")
    export.add_argument("--comparison",help="Optional measured comparison JSON bound to these weights")
    args = p.parse_args()
    if args.command == "prepare":
        from .data import prepare
        prepare(args)
    elif args.command == "train":
        from .train import train as run
        run(args)
    elif args.command == "benchmark":
        from .benchmark import run
        run(args)
    elif args.command == "serve":
        from .demo import serve as run
        run(args)
    elif args.command == "export":
        from .export import export as run
        run(args)
    else:
        import json
        import torch
        from .inference import load, continue_text, infill
        torch.set_num_threads(4)
        model, tokenizer = load(args.model, args.device)
        options = dict(steps=args.steps, temperature=args.temperature, seed=args.seed)
        if args.suffix is None:
            from pathlib import Path
            info = json.loads((Path(args.model)/"training_info.json").read_text())
            use_causal = args.generation_mode == "causal" or (args.generation_mode == "auto" and info.get("continuation_mode") == "autoregressive")
            if use_causal:
                if info.get("continuation_mode") != "autoregressive":
                    raise ValueError("This checkpoint was not trained for causal generation")
                from .hybrid import generate
                result = generate(model, tokenizer, args.prompt, max_new_tokens=args.new_tokens,
                                  finish_sentence_tokens=args.finish_sentence_tokens, **options)
            else:
                result = continue_text(model, tokenizer, args.prompt, max_new_tokens=args.new_tokens, block_size=args.block_size, **options)
        else:
            result = infill(model, tokenizer, args.prompt, args.suffix, missing_tokens=args.missing_tokens, **options)
        print(json.dumps(result, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()