# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 Lucasgitguy-co
import argparse
import importlib.util
import json
import platform
import shutil
import sys
from pathlib import Path

import torch


def parser():
    p = argparse.ArgumentParser(description="Local image-to-Gaussian student research prototype")
    sub = p.add_subparsers(dest="command", required=True)
    sub.add_parser("doctor", help="Check dependencies, Metal availability and disk space; no downloads")
    demo = sub.add_parser("demo-data", help="Create synthetic code-test fixtures, NOT teacher data")
    demo.add_argument("--out", required=True)
    demo.add_argument("--count", type=int, default=12)
    demo.add_argument("--grid", type=int, choices=[8,16,32], default=8)
    cache = sub.add_parser("cache", help="Cache Idol PLY with optional aligned geometry teacher meshes")
    cache.add_argument("--image", required=True, help="Exact prepared/framed input used by teacher")
    cache.add_argument("--idol-ply", required=True)
    cache.add_argument("--out", required=True)
    cache.add_argument("--id", required=True)
    cache.add_argument("--object-id", help="Shared ID for all images of the same object, to prevent validation leakage")
    cache.add_argument("--grid", type=int, choices=[8,16,32], default=16)
    cache.add_argument("--coarse-mesh", help="InstantMesh output; this command imports, does not run InstantMesh")
    cache.add_argument("--coarse-transform", help="JSON 4x4 mesh-to-Idol-PLY matrix")
    cache.add_argument("--fine-mesh", help="Hi3D output; this command imports, does not run Hi3D")
    cache.add_argument("--fine-transform", help="JSON 4x4 mesh-to-Idol-PLY matrix")
    train = sub.add_parser("train", help="Train the student on cached targets")
    train.add_argument("--manifest", required=True)
    train.add_argument("--out", required=True)
    train.add_argument("--device", choices=["auto","mps","cpu"], default="auto")
    train.add_argument("--steps", type=int, default=1000)
    train.add_argument("--batch-size", type=int, default=1)
    train.add_argument("--width", type=int, default=128)
    train.add_argument("--lr", type=float, default=3e-4)
    train.add_argument("--seed", type=int, default=42)
    train.add_argument("--eval-every", type=int, default=50)
    train.add_argument("--validation-fraction", type=float, default=.2)
    train.add_argument("--allow-synthetic", action="store_true")
    generate = sub.add_parser("generate", help="Generate standard 3DGS PLY with a trained student")
    generate.add_argument("--checkpoint", required=True)
    generate.add_argument("--image", required=True, help="Centered cutout or teacher-prepared image, black background")
    generate.add_argument("--out", required=True)
    generate.add_argument("--device", choices=["auto","mps","cpu"], default="auto")
    generate.add_argument("--threshold", type=float, default=.5)
    generate.add_argument("--max-splats", type=int, default=32768)
    render = sub.add_parser("preview", help="CPU diagnostic turntable contact sheet")
    render.add_argument("--ply", required=True)
    render.add_argument("--out", required=True)
    teacher = sub.add_parser("teacher-idol", help="EXPERIMENTAL staged Mac runner; weights not bundled")
    teacher.add_argument("--image", required=True)
    teacher.add_argument("--out", required=True)
    teacher.add_argument("--cache", default=".cache/huggingface")
    teacher.add_argument("--weights", default="speridlabs/idol")
    teacher.add_argument("--device", choices=["mps","cpu"], default="mps")
    teacher.add_argument("--mode", choices=["real","synthetic"], default="synthetic")
    teacher.add_argument("--steps", type=int, default=20)
    teacher.add_argument("--gaussians", type=int, default=32768)
    teacher.add_argument("--query-chunk", type=int, default=128)
    teacher.add_argument("--seed", type=int, default=42)
    teacher.add_argument("--allow-download", action="store_true", help="Explicitly allow downloading model weights and gated encoders")
    return p


def main():
    p = parser()
    args = p.parse_args()
    try:
        if args.command == "doctor":
            report = {"python": platform.python_version(), "platform": platform.platform(), "torch": str(torch.__version__),
                      "mps_available": torch.backends.mps.is_available(), "mps_built": torch.backends.mps.is_built(),
                      "free_disk_gib": round(shutil.disk_usage(Path.cwd()).free / 1024**3, 2),
                      "idol_installed": importlib.util.find_spec("idol") is not None,
                      "teacher_status": {"Idol": "experimental adapter, real weights unverified",
                                         "InstantMesh": "cached mesh import only; no Mac execution adapter",
                                         "Hi3D": "cached mesh import only; no Mac execution adapter"},
                      "weights_bundled": False, "quantization": "FP32 student; no binary kernels or 1-bit claims"}
            print(json.dumps(report, indent=2))
        elif args.command == "demo-data":
            from .demo import create_demo
            print(create_demo(args.out, args.count, args.grid))
        elif args.command == "cache":
            from .data import cache_sample
            from .splats import read_ply
            print(json.dumps(cache_sample(args.out, args.id, args.image, read_ply(args.idol_ply), args.grid,
                coarse=args.coarse_mesh, coarse_transform=args.coarse_transform,
                fine=args.fine_mesh, fine_transform=args.fine_transform,
                extra={"object_id": args.object_id or args.id, "idol_ply": str(Path(args.idol_ply).resolve())}), indent=2))
        elif args.command == "train":
            from .engine import train
            train(args)
        elif args.command == "generate":
            from .engine import generate
            print(json.dumps(generate(args.checkpoint, args.image, args.out, args.device, args.threshold, args.max_splats), indent=2))
        elif args.command == "preview":
            from PIL import Image, ImageDraw
            import numpy as np
            from .splats import preview, read_ply
            splats = read_ply(args.ply)
            sheet = Image.new("RGB", (1024, 280), "#111111")
            draw = ImageDraw.Draw(sheet)
            for i in range(4):
                sheet.paste(preview(splats, yaw=i*np.pi/2), (i*256, 24))
                draw.text((i*256+8, 6), f"{i*90} degrees - diagnostic view", fill="white")
            out = Path(args.out)
            if out.exists():
                raise ValueError("Preview already exists")
            out.parent.mkdir(parents=True, exist_ok=True)
            sheet.save(out)
        elif args.command == "teacher-idol":
            from .teacher_idol import run_idol
            run_idol(args)
    except (ValueError, RuntimeError, FileNotFoundError, ImportError) as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 2
    return 0


if __name__ == "__main__":
    raise SystemExit(main())