import argparse
import json
from pathlib import Path


def main():
    parser = argparse.ArgumentParser(description="Private Solomon BF16 MLX tooling")
    commands = parser.add_subparsers(dest="command", required=True)
    verify = commands.add_parser("verify-source")
    verify.add_argument("directory")
    download = commands.add_parser("download-base")
    download.add_argument("--output", default="snapshots/base")
    convert = commands.add_parser("prepare")
    convert.add_argument("--base", default="snapshots/base")
    convert.add_argument("--solomon", default="snapshots/solomon")
    convert.add_argument("--manifest", default="snapshots/base-manifest.json")
    convert.add_argument("--output", default="models/quality")
    decide = commands.add_parser("decide")
    decide.add_argument("--model", default="models/quality")
    decide.add_argument("--document", required=True)
    decide.add_argument("--questions", required=True)
    decide.add_argument(
        "--evidence", choices=["none", "support", "sufficiency", "removal"], default="support"
    )
    args = parser.parse_args()
    if args.command == "verify-source":
        from .artifacts import verify_release

        print(json.dumps(verify_release(args.directory), indent=2))
    elif args.command == "download-base":
        from huggingface_hub import HfApi, snapshot_download

        from .artifacts import BASE_REVISION
        from .prepare import verify_base

        model = HfApi().model_info("Qwen/Qwen3.8-27B", revision=BASE_REVISION, files_metadata=True)
        manifest = {
            "revision": model.sha,
            "files": [
                {
                    "name": f.rfilename,
                    "size": f.size,
                    "blob_id": f.blob_id,
                    "sha256": f.lfs.sha256 if f.lfs else None,
                }
                for f in model.siblings
            ],
        }
        Path(args.output).parent.mkdir(parents=True, exist_ok=True)
        Path(args.output + "-manifest.json").write_text(json.dumps(manifest, indent=2))
        snapshot_download("Qwen/Qwen3.8-27B", revision=BASE_REVISION, local_dir=args.output, max_workers=18)
        verified = verify_base(args.output, manifest)
        Path(args.output + "-verified.json").write_text(json.dumps(verified, indent=2))
    elif args.command == "prepare":
        from .prepare import prepare

        prepare(args.base, args.solomon, args.output, args.manifest)
    else:
        from .api import Solomon

        model = Solomon.load(args.model)
        with model.prefill(Path(args.document).read_text()) as state:
            print(
                json.dumps(
                    model.decide(
                        state=state,
                        questions=json.loads(Path(args.questions).read_text()),
                        evidence=args.evidence,
                    ),
                    indent=2,
                )
            )


if __name__ == "__main__":
    main()