from __future__ import annotations

import argparse
import json
from collections.abc import Sequence
from pathlib import Path

from barunlm.evaluation.generation import generate_manifest
from barunlm.evaluation.mobile_actions import write_scores


def _is_sha256(value: str) -> bool:
    return len(value) == 64 and all(character in "0123456789abcdef" for character in value)


def _hash_map(value: str | None) -> dict[str, str] | None:
    if value is None:
        return None
    payload = json.loads(Path(value).read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise TypeError("checkpoint hash file must contain an object")
    if "file_sha256" in payload:
        payload = payload["file_sha256"]
    if not isinstance(payload, dict):
        raise TypeError("checkpoint file_sha256 must contain an object")
    return {str(name): str(digest) for name, digest in payload.items()}


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Generate and strictly score a Mobile Actions development manifest."
    )
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--checkpoint-format", choices=("float", "int8"), default="float")
    parser.add_argument("--checkpoint-hashes", help="JSON hash map for a base checkpoint")
    parser.add_argument(
        "--int8-manifest-sha256",
        help="Required out-of-band manifest SHA-256 for an int8 checkpoint.",
    )
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--manifest-sha256", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--device", choices=("cpu", "cuda"), required=True)
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--max-new-tokens", type=int, default=192)
    return parser


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = build_parser()
    args = parser.parse_args(argv)
    if args.checkpoint_format == "float" and args.int8_manifest_sha256 is not None:
        parser.error("--int8-manifest-sha256 is valid only with --checkpoint-format int8")
    if args.checkpoint_format == "int8":
        if args.checkpoint_hashes is not None:
            parser.error("--checkpoint-hashes is valid only with --checkpoint-format float")
        if args.int8_manifest_sha256 is None:
            parser.error("--checkpoint-format int8 requires --int8-manifest-sha256")
        if not _is_sha256(args.int8_manifest_sha256):
            parser.error("--int8-manifest-sha256 must be a lowercase SHA-256")
        if args.device != "cpu":
            parser.error("--checkpoint-format int8 requires --device cpu")
    return args


def main(argv: Sequence[str] | None = None) -> int:
    args = parse_args(argv)
    checkpoint_hashes = _hash_map(args.checkpoint_hashes)
    args.output.mkdir(parents=True, exist_ok=False)
    predictions = args.output / "predictions.jsonl"
    generation = generate_manifest(
        checkpoint_dir=args.checkpoint,
        manifest_path=args.manifest,
        manifest_sha256=args.manifest_sha256,
        predictions_path=predictions,
        device_name=args.device,
        batch_size=args.batch_size,
        max_new_tokens=args.max_new_tokens,
        expected_checkpoint_sha256=checkpoint_hashes,
        checkpoint_format=args.checkpoint_format,
        expected_int8_manifest_sha256=args.int8_manifest_sha256,
    )
    score_paths = write_scores(args.manifest, predictions, args.output / "scores")
    print(
        json.dumps(
            {
                "generation": generation.to_dict(),
                "scores": {name: str(path) for name, path in score_paths.items()},
            },
            allow_nan=False,
            sort_keys=True,
        )
    )
    return 0


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