#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import shutil
from pathlib import Path


CONFIGS = [
    "A0_full",
    "A1_no_planner",
    "A2_no_neo4j",
    "A3_no_vector",
    "A4_one_hop",
    "A5_no_bge",
]


PROJECT_ROOT = Path(__file__).resolve().parents[1]


def resolve_path(value: str | Path) -> Path:
    path = Path(value)
    if not path.is_absolute():
        path = PROJECT_ROOT / path
    return path


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--base-dir",
        default="experiments/rag_ablation_20260624",
        help="Base ablation experiment directory, relative to the project root unless absolute.",
    )
    parser.add_argument("--first-20-dir", default="")
    parser.add_argument("--remaining-80-dir", default="")
    parser.add_argument("--output-dir", default="")
    parser.add_argument("--questions", default="")
    parser.add_argument(
        "--allow-existing-output",
        action="store_true",
        help="Allow writing into an existing output directory without deleting it.",
    )
    args = parser.parse_args()

    base = resolve_path(args.base_dir)
    first_20 = resolve_path(args.first_20_dir) if args.first_20_dir else base / "paper_module_strict_official_20q_7907_20260624"
    remaining_80 = resolve_path(args.remaining_80_dir) if args.remaining_80_dir else base / "paper_module_remaining_80q_7907_20260624"
    output = resolve_path(args.output_dir) if args.output_dir else base / "paper_module_strict_official_100q_7907_20260624"
    questions_path = (
        resolve_path(args.questions)
        if args.questions
        else base / "00_inputs" / "QA100_questions_gold_answers_source_text_no_srt_audited_v1_20260624.json"
    )
    questions = json.loads(questions_path.read_text(encoding="utf-8-sig"))
    ids = [str(item["id"]) for item in questions]
    first_ids = {
        str(item["id"])
        for item in json.loads(
            (first_20 / "results" / "selected_questions.json").read_text(
                encoding="utf-8"
            )
        )
    }
    remaining_ids = {
        str(item["id"])
        for item in json.loads(
            (remaining_80 / "results" / "selected_questions.json").read_text(
                encoding="utf-8"
            )
        )
    }
    if first_ids & remaining_ids:
        raise RuntimeError(f"Overlapping IDs: {sorted(first_ids & remaining_ids)}")
    if first_ids | remaining_ids != set(ids):
        missing = set(ids) - (first_ids | remaining_ids)
        extra = (first_ids | remaining_ids) - set(ids)
        raise RuntimeError(f"Question mismatch missing={sorted(missing)} extra={sorted(extra)}")

    if output.exists() and not args.allow_existing_output:
        raise FileExistsError(
            f"Refusing to write into existing output directory: {output}. "
            "Choose a new --output-dir or pass --allow-existing-output."
        )
    (output / "results").mkdir(parents=True, exist_ok=args.allow_existing_output)
    (output / "analysis").mkdir(parents=True, exist_ok=args.allow_existing_output)
    (output / "judge").mkdir(parents=True, exist_ok=args.allow_existing_output)
    (output / "results" / "selected_questions.json").write_text(
        json.dumps(questions, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )

    counts: dict[str, int] = {}
    for config in CONFIGS:
        target_dir = output / "results" / config
        target_dir.mkdir(parents=True, exist_ok=args.allow_existing_output)
        copied = 0
        for qid in ids:
            source_root = first_20 if qid in first_ids else remaining_80
            source = source_root / "results" / config / f"{qid}.json"
            if not source.exists():
                raise FileNotFoundError(source)
            record = json.loads(source.read_text(encoding="utf-8"))
            if record.get("error") or not record.get("response"):
                raise RuntimeError(f"Invalid result: {source}")
            shutil.copy2(source, target_dir / source.name)
            copied += 1
        counts[config] = copied
    manifest = {
        "questions": len(ids),
        "configs": CONFIGS,
        "expected_results": len(ids) * len(CONFIGS),
        "copied_results": sum(counts.values()),
        "per_config": counts,
        "sources": {
            "first_20": str(first_20),
            "remaining_80": str(remaining_80),
            "gold_questions": str(questions_path),
        },
    }
    (output / "merge_manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(manifest, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()