from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

from .dataset import (
    DEFAULT_GROUP_FIELDS,
    assess_model_promotion,
    assign_splits,
    matrix_from_records,
    records_manifest,
    split_by_group,
    split_receipt,
)
from .demo import generate_records, generate_synthetic_fixture, run_demo
from .inference import infer_capture
from .io import load_capture, read_json, write_json
from .model import build_model_card, evaluate_binary, fit_logistic, load_model, save_model
from .quality import quality_check_capture


def _load_records(path: str | Path) -> list[dict[str, Any]]:
    value = read_json(path)
    records = value.get("records")
    if not isinstance(records, list) or not all(isinstance(record, dict) for record in records):
        raise ValueError("Records JSON must contain an object array named records")
    return records


def parser() -> argparse.ArgumentParser:
    root = argparse.ArgumentParser(
        prog="m3m-disease",
        description="Research-only RGB and Mavic 3M optical disease-screening engineering prototype",
    )
    commands = root.add_subparsers(dest="command", required=True)

    fixture = commands.add_parser("fixture", help="Create deterministic synthetic RGB and G/R/RE/NIR fixtures")
    fixture.add_argument("--output", required=True)
    fixture.add_argument("--seed", type=int, default=20260804)

    demo = commands.add_parser("demo", help="Run fixture, train, evaluation and inference end to end")
    demo.add_argument("--output", required=True)
    demo.add_argument("--seed", type=int, default=20260804)

    validate = commands.add_parser("validate", help="Run the fail-closed capture QC gate")
    validate.add_argument("--manifest", required=True)
    validate.add_argument("--output", required=True)

    train = commands.add_parser("train", help="Train and serialize the auditable L2 logistic model")
    train.add_argument("--records", required=True)
    train.add_argument("--output", required=True)
    train.add_argument("--seed", type=int, default=20260804)

    evaluate = commands.add_parser("evaluate", help="Evaluate a JSON model on one presplit dataset partition")
    evaluate.add_argument("--records", required=True)
    evaluate.add_argument("--model", required=True)
    evaluate.add_argument("--split", choices=["train", "validation", "test"], default="test")
    evaluate.add_argument("--output", required=True)

    predict = commands.add_parser("predict", help="Screen one registered capture and emit an inference receipt")
    predict.add_argument("--manifest", required=True)
    predict.add_argument("--model", required=True)
    predict.add_argument("--patch-size", type=int, default=32)
    predict.add_argument("--output", required=True)
    return root


def main(argv: list[str] | None = None) -> int:
    args = parser().parse_args(argv)
    if args.command == "fixture":
        output = Path(args.output)
        manifests = generate_synthetic_fixture(output, seed=args.seed)
        records = generate_records(manifests)
        write_json(output / "dataset-records.json", {"schemaVersion": 1, "records": records})
        write_json(output / "dataset-manifest.json", records_manifest(records))
        result: dict[str, Any] = {
            "status": "PASSED",
            "dataProvenance": "synthetic",
            "captureCount": len(manifests),
            "sampleCount": len(records),
            "fieldPerformanceEstablished": False,
        }
    elif args.command == "demo":
        result = run_demo(args.output, seed=args.seed)
    elif args.command == "validate":
        rgb, bands, metadata = load_capture(args.manifest)
        result = quality_check_capture(rgb, bands, metadata)
        write_json(args.output, result)
    elif args.command == "train":
        output = Path(args.output)
        records = _load_records(args.records)
        assignments = split_by_group(records, seed=args.seed, group_fields=DEFAULT_GROUP_FIELDS)
        assigned = assign_splits(records, assignments, group_fields=DEFAULT_GROUP_FIELDS)
        receipt = split_receipt(records, assignments, group_fields=DEFAULT_GROUP_FIELDS, seed=args.seed)
        train_x, train_y, _ = matrix_from_records(assigned, split="train")
        test_x, test_y, _ = matrix_from_records(assigned, split="test")
        promotion = assess_model_promotion(assigned)
        provenance = "field" if all(record["data_provenance"] == "field" for record in records) else "synthetic"
        model = fit_logistic(
            train_x,
            train_y,
            data_provenance=provenance,
            promotion_receipt=promotion,
        )
        model_path = save_model(output / "model.json", model)
        evaluation = evaluate_binary(model, test_x, test_y)
        card = build_model_card(model, split_receipt=receipt, evaluation=evaluation, promotion=promotion)
        write_json(output / "model-card.json", card)
        write_json(output / "split-receipt.json", receipt)
        write_json(output / "evaluation.json", evaluation)
        write_json(output / "dataset-records-presplit.json", {"schemaVersion": 1, "records": assigned})
        result = {
            "status": "PASSED",
            "model": str(model_path),
            "promotionStatus": promotion["status"],
            "fieldPerformanceEstablished": False,
        }
    elif args.command == "evaluate":
        records = _load_records(args.records)
        matrix, labels, _ = matrix_from_records(records, split=args.split)
        result = evaluate_binary(load_model(args.model), matrix, labels)
        write_json(args.output, result)
    else:
        result = infer_capture(load_model(args.model), args.manifest, patch_size=args.patch_size, output_path=args.output)
    print(json.dumps(result, ensure_ascii=False, sort_keys=True))
    return 0
