from __future__ import annotations

import argparse, json
from pathlib import Path
from .describe import describe_samples
from .pipeline import load_samples, run_stage, scan, select_samples
from .regression import regression_sample_ids, summarize_regression, write_regression_manifest
from .reports import generate_markdown_report, generate_reports, read_json

DEFAULT_INPUT = Path("data/cadfs-sample/CADFS_test")
DEFAULT_OUTPUT = Path("cadfs_to_cdsl/output")


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Convert CADFS FeatureScript to CDSL and validate against STEP")
    commands = parser.add_subparsers(dest="command", required=True)
    for name in ("scan", "convert", "rebuild", "compare", "report", "pipeline", "describe", "regression-select", "regression"):
        command = commands.add_parser(name)
        command.add_argument("--input", type=Path, default=DEFAULT_INPUT)
        command.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
        if name == "regression-select":
            command.add_argument("--manifest", type=Path, default=Path("cadfs_to_cdsl/regression/manifest.json"))
        elif name == "regression":
            command.add_argument("--manifest", type=Path, default=Path("cadfs_to_cdsl/regression/manifest.json"))
            command.add_argument("--tier", choices=("engine", "conversion", "all"), default="engine")
            command.add_argument("--stage", choices=("convert", "rebuild", "compare", "pipeline"), default="rebuild")
            command.add_argument("--workers", type=int, default=1)
            command.add_argument("--compare-mode", choices=("rp", "strict"), default="rp")
            command.add_argument("--timeout-seconds", type=float, default=30.0, help="per-model OCC timeout (default: 30)")
            command.add_argument("--resume", action="store_true", help="reuse cached per-sample stage results instead of rerunning them")
        elif name == "describe":
            command.add_argument("--shard", help="input shard directory to describe, for example 0005")
            command.add_argument("--sample-id", action="append")
            command.add_argument("--offset", type=int, default=0)
            command.add_argument("--limit", type=int)
            command.add_argument("--seed", type=int)
            command.add_argument("--workers", type=int, default=1)
            command.add_argument("--mode", choices=("local", "hybrid", "vision"), default="hybrid")
            command.add_argument("--force", action="store_true")
        elif name != "scan":
            command.add_argument("--sample-id", action="append")
            command.add_argument("--offset", type=int, default=0)
            command.add_argument("--limit", type=int)
            command.add_argument("--seed", type=int)
            command.add_argument("--workers", type=int, default=1)
            command.add_argument("--compare-mode", choices=("rp", "strict"), default="rp")
            command.add_argument("--force", action="store_true")
            command.add_argument("--timeout-seconds", type=float, default=30.0, help="per-model OCC timeout (default: 30)")
        if name == "report":
            command.add_argument("--markdown", action="store_true", help="also write full_run_report.md")
    return parser


def main(argv: list[str] | None = None) -> int:
    args = _parser().parse_args(argv)
    if args.command != "describe":
        args.output.mkdir(parents=True, exist_ok=True)
    if args.command == "scan":
        records = scan(args.input, args.output); result = {"sample_count": len(records), "output": str(args.output / "dataset_index.json")}
    elif args.command == "regression-select":
        manifest = write_regression_manifest(args.output, args.manifest)
        result = {"manifest": str(args.manifest), "selected_sample_count": manifest["selected_sample_count"], "engine_sample_count": manifest["engine_sample_count"]}
    elif args.command == "regression":
        if not 1 <= args.workers <= 8: raise ValueError("--workers must be between 1 and 8")
        if args.timeout_seconds <= 0: raise ValueError("--timeout-seconds must be positive")
        sample_ids = regression_sample_ids(args.manifest, args.tier)
        samples = select_samples(load_samples(args.input, args.output), sample_ids=sample_ids)
        records = run_stage(args.stage, samples, args.output, force=not args.resume, compare_mode=args.compare_mode, timeout_seconds=args.timeout_seconds, workers=args.workers)
        result = {"tier": args.tier, "stage": args.stage, "manifest": str(args.manifest), **summarize_regression(records)}
    elif args.command == "describe":
        records, result = describe_samples(
            args.input,
            shard=args.shard,
            sample_ids=args.sample_id,
            mode=args.mode,
            offset=args.offset,
            limit=args.limit,
            seed=args.seed,
            force=args.force,
            workers=args.workers,
        )
    elif args.command == "report":
        manifest = args.output / "manifest.jsonl"; records = [json.loads(line) for line in manifest.read_text().splitlines() if line.strip()]
        result = generate_reports(args.output, records)
        if args.markdown:
            command_text = "python -m cadfs_to_cdsl report --markdown"
            result = {**result, "markdown_report": str(generate_markdown_report(args.output, records, input_root=args.input, command=command_text))}
    else:
        if not 1 <= args.workers <= 8: raise ValueError("--workers must be between 1 and 8")
        samples = select_samples(load_samples(args.input, args.output), sample_ids=args.sample_id, offset=args.offset, limit=args.limit, seed=args.seed)
        if args.timeout_seconds <= 0: raise ValueError("--timeout-seconds must be positive")
        records = run_stage(args.command, samples, args.output, force=args.force, compare_mode=args.compare_mode, timeout_seconds=args.timeout_seconds, workers=args.workers)
        counts: dict[str, int] = {}
        for record in records: counts[record["status"]] = counts.get(record["status"], 0) + 1
        result = {"sample_count": len(records), "statuses": counts, "summary": str(args.output / "summary.json")}
    print(json.dumps(result, ensure_ascii=True, indent=2, sort_keys=True)); return 0