#!/usr/bin/env python3

import argparse
import json
from pathlib import Path
from typing import Any, Dict, Iterable, List, Tuple


FIRST_PASS_GATES = {
    "prereq_op_selector",
    "target_file_selection",
    "acceptance_command_selection",
    "required_tool_selection",
}
REPO_ROOT = Path(__file__).resolve().parent.parent


def load_json(path: Path) -> Dict[str, Any]:
    with path.open("r", encoding="utf-8") as fh:
        return json.load(fh)


def iter_run_dirs(root: Path) -> Iterable[Path]:
    for child in sorted(root.iterdir()):
        if not child.is_dir():
            continue
        if (child / "01_intake.json").exists() and (child / "02_generate_taskitems.json").exists():
            yield child


def nonempty_list(value: Any) -> List[Any]:
    if isinstance(value, list):
        return value
    return []


def task_context(intake: Dict[str, Any], task: Dict[str, Any]) -> Dict[str, Any]:
    return {
        "normalized_requirements": intake.get("normalizedRequirements", []),
        "parsed_spec": intake.get("parsedSpec", {}),
        "conflict_signals": intake.get("conflictSignals", []),
        "ambiguity_signals": intake.get("ambiguitySignals", []),
        "task_id": task.get("taskId"),
        "title": task.get("title"),
        "milestone_id": task.get("milestoneId"),
        "dependency_task_ids": nonempty_list(task.get("dependencyTaskIds")),
        "confidence": task.get("confidence"),
        "ambiguity_count": task.get("ambiguityCount"),
        "reasons": nonempty_list(task.get("reasons")),
    }


def build_sibling_maps(tasks: List[Dict[str, Any]]) -> Dict[str, Dict[str, int]]:
    target_counts: Dict[str, int] = {}
    acceptance_counts: Dict[str, int] = {}
    tool_counts: Dict[str, int] = {}
    for task in tasks:
        contract = task.get("executionContract", {})
        targets = tuple(nonempty_list(contract.get("targetFiles")))
        acceptance = tuple(nonempty_list(contract.get("acceptanceCommands")))
        tools = tuple(nonempty_list(contract.get("requiredTools")))
        if targets:
            target_counts[json.dumps(targets)] = target_counts.get(json.dumps(targets), 0) + 1
        if acceptance:
            acceptance_counts[json.dumps(acceptance)] = acceptance_counts.get(json.dumps(acceptance), 0) + 1
        if tools:
            tool_counts[json.dumps(tools)] = tool_counts.get(json.dumps(tools), 0) + 1
    return {
        "targets": target_counts,
        "acceptance": acceptance_counts,
        "tools": tool_counts,
    }


def base_row(
    gate_id: str,
    run_dir: Path,
    task: Dict[str, Any],
    contract_path: Path,
    provenance: str,
) -> Dict[str, Any]:
    return {
        "gate_id": gate_id,
        "run_id": run_dir.name,
        "task_id": task.get("taskId"),
        "decision_context": {},
        "decision_label": {},
        "decision_provenance": provenance,
        "source_artifacts": [
            str(run_dir / "01_intake.json"),
            str(run_dir / "02_generate_taskitems.json"),
            str(contract_path),
        ],
        "rejection_flags": [],
    }


def extract_prereq_row(
    run_dir: Path,
    intake: Dict[str, Any],
    task: Dict[str, Any],
    contract_path: Path,
) -> Dict[str, Any]:
    row = base_row("prereq_op_selector", run_dir, task, contract_path, "direct_bundled")
    prereqs = nonempty_list(task.get("prerequisiteOps"))
    accepted = {"validate-intake", "architect-review", "resolve-dependencies"}
    invalid = sorted(op for op in prereqs if op not in accepted)
    row["decision_context"] = task_context(intake, task)
    row["decision_label"] = {
        "needs_validate_intake": "validate-intake" in prereqs,
        "needs_architect_review": "architect-review" in prereqs,
        "needs_resolve_dependencies": "resolve-dependencies" in prereqs,
        "source_values": prereqs,
    }
    if invalid:
        row["rejection_flags"].append(f"invalid_prerequisite_values:{','.join(invalid)}")
    if not row["decision_context"]["title"]:
        row["rejection_flags"].append("missing_task_title")
    if row["decision_context"]["confidence"] is None:
        row["rejection_flags"].append("missing_confidence")
    return row


def extract_target_files_row(
    run_dir: Path,
    intake: Dict[str, Any],
    task: Dict[str, Any],
    contract_path: Path,
    sibling_maps: Dict[str, Dict[str, int]],
    sibling_count: int,
) -> Dict[str, Any]:
    row = base_row("target_file_selection", run_dir, task, contract_path, "direct_bundled")
    execution_contract = task.get("executionContract", {})
    targets = nonempty_list(execution_contract.get("targetFiles"))
    row["decision_context"] = task_context(intake, task)
    row["decision_label"] = {
        "target_files": targets,
    }
    if not targets:
        row["rejection_flags"].append("empty_target_files")
    if sibling_count > 1 and targets:
        key = json.dumps(tuple(targets))
        if sibling_maps["targets"].get(key, 0) == sibling_count:
            row["rejection_flags"].append("sibling_uniform_target_files")
    return row


def extract_acceptance_row(
    run_dir: Path,
    intake: Dict[str, Any],
    task: Dict[str, Any],
    contract_path: Path,
    sibling_maps: Dict[str, Dict[str, int]],
    sibling_count: int,
) -> Dict[str, Any]:
    row = base_row("acceptance_command_selection", run_dir, task, contract_path, "direct_bundled")
    execution_contract = task.get("executionContract", {})
    commands = nonempty_list(execution_contract.get("acceptanceCommands"))
    row["decision_context"] = task_context(intake, task)
    row["decision_context"]["target_files"] = nonempty_list(execution_contract.get("targetFiles"))
    row["decision_label"] = {
        "acceptance_commands": commands,
    }
    if not commands:
        row["rejection_flags"].append("empty_acceptance_commands")
    if not row["decision_context"]["target_files"]:
        row["rejection_flags"].append("missing_target_files_context")
    if sibling_count > 1 and commands:
        key = json.dumps(tuple(commands))
        if sibling_maps["acceptance"].get(key, 0) == sibling_count:
            row["rejection_flags"].append("sibling_uniform_acceptance_commands")
    return row


def extract_required_tools_row(
    run_dir: Path,
    intake: Dict[str, Any],
    task: Dict[str, Any],
    contract_path: Path,
    sibling_maps: Dict[str, Dict[str, int]],
    sibling_count: int,
) -> Dict[str, Any]:
    row = base_row("required_tool_selection", run_dir, task, contract_path, "direct_bundled")
    execution_contract = task.get("executionContract", {})
    tools = nonempty_list(execution_contract.get("requiredTools"))
    row["decision_context"] = task_context(intake, task)
    row["decision_context"]["prerequisite_ops"] = nonempty_list(task.get("prerequisiteOps"))
    row["decision_context"]["target_files"] = nonempty_list(execution_contract.get("targetFiles"))
    row["decision_label"] = {
        "required_tools": tools,
    }
    if not tools:
        row["rejection_flags"].append("empty_required_tools")
    if sibling_count > 1 and tools:
        key = json.dumps(tuple(tools))
        if sibling_maps["tools"].get(key, 0) == sibling_count:
            row["rejection_flags"].append("sibling_uniform_required_tools")
    return row


def extract_rows_from_run(run_dir: Path, contract_path: Path) -> Tuple[List[Dict[str, Any]], Dict[str, int]]:
    intake = load_json(run_dir / "01_intake.json")
    generated = load_json(run_dir / "02_generate_taskitems.json")
    tasks = nonempty_list(generated.get("tasks"))
    sibling_maps = build_sibling_maps(tasks)
    rows: List[Dict[str, Any]] = []
    stats = {"runs": 1, "tasks": len(tasks), "rows": 0, "rejected_rows": 0}
    sibling_count = len(tasks)
    for task in tasks:
        gate_rows = [
            extract_prereq_row(run_dir, intake, task, contract_path),
            extract_target_files_row(run_dir, intake, task, contract_path, sibling_maps, sibling_count),
            extract_acceptance_row(run_dir, intake, task, contract_path, sibling_maps, sibling_count),
            extract_required_tools_row(run_dir, intake, task, contract_path, sibling_maps, sibling_count),
        ]
        for row in gate_rows:
            rows.append(row)
            stats["rows"] += 1
            if row["rejection_flags"]:
                stats["rejected_rows"] += 1
    return rows, stats


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Extract RSA gate rows from whetstone_DSL taskitem runs.")
    parser.add_argument(
        "--contracts",
        default="semantic/rsa_gate_extraction_contracts_v0.json",
        help="Path to the extraction contract JSON.",
    )
    parser.add_argument(
        "--taskitem-runs-root",
        required=True,
        help="Root directory containing whetstone_DSL taskitem run folders.",
    )
    parser.add_argument(
        "--output",
        required=True,
        help="Output NDJSON path for extracted rows.",
    )
    parser.add_argument(
        "--limit-runs",
        type=int,
        default=0,
        help="Optional limit on how many run directories to process.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    contract_path = Path(args.contracts)
    if not contract_path.is_absolute():
        contract_path = (REPO_ROOT / contract_path).resolve()
    contracts = load_json(contract_path)
    missing = FIRST_PASS_GATES.difference(contracts.get("contracts", {}).keys())
    if missing:
        raise SystemExit(f"missing first-pass gate contracts: {sorted(missing)}")

    runs_root = Path(args.taskitem_runs_root).resolve()
    output_path = Path(args.output).resolve()
    output_path.parent.mkdir(parents=True, exist_ok=True)

    rows: List[Dict[str, Any]] = []
    totals = {"runs": 0, "tasks": 0, "rows": 0, "rejected_rows": 0}
    for index, run_dir in enumerate(iter_run_dirs(runs_root), start=1):
        if args.limit_runs and index > args.limit_runs:
            break
        run_rows, run_stats = extract_rows_from_run(run_dir, contract_path)
        rows.extend(run_rows)
        for key, value in run_stats.items():
            totals[key] += value

    with output_path.open("w", encoding="utf-8") as fh:
        for row in rows:
            fh.write(json.dumps(row, sort_keys=True) + "\n")

    summary = {
        "contracts_path": str(contract_path),
        "runs_root": str(runs_root),
        "output_path": str(output_path),
        "runs_processed": totals["runs"],
        "tasks_processed": totals["tasks"],
        "rows_emitted": totals["rows"],
        "rows_with_rejection_flags": totals["rejected_rows"],
    }
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0


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