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

import json
from collections import Counter
from pathlib import Path
from typing import Any


IN_FILE = Path("runs/subproblem_decomposition/random20_decomposed.json")
OUT_FILE = Path("runs/subproblem_decomposition/random20_coarse_decomposed.json")
REPORT_FILE = Path("runs/subproblem_decomposition/coarse_report.json")


COARSE_MAP = {
    "understand_task": "understand",
    "explore_repo": "locate",
    "locate_relevant_code": "locate",
    "inspect_code": "diagnose",
    "reproduce_or_probe": "diagnose",
    "edit_solution": "fix",
    "verify_solution": "verify",
    "cleanup": "finalize",
    "review_diff": "finalize",
    "submit": "finalize",
    "other": "diagnose",
}


COARSE_GOALS = {
    "understand": "Understand the issue, repository context, constraints, and success criteria.",
    "locate": "Find the files, modules, symbols, or tests likely responsible for the issue.",
    "diagnose": "Inspect behavior and evidence to determine the concrete root cause.",
    "fix": "Edit the implementation or tests to address the root cause.",
    "verify": "Run targeted checks to confirm the fix and catch regressions.",
    "finalize": "Clean temporary artifacts, inspect the final diff, and submit the solution.",
}


def dedupe(items: list[Any], limit: int | None = None) -> list[Any]:
    seen = set()
    out = []
    for item in items:
        key = json.dumps(item, ensure_ascii=False, sort_keys=True) if isinstance(item, (dict, list)) else str(item)
        if key in seen:
            continue
        seen.add(key)
        out.append(item)
        if limit and len(out) >= limit:
            break
    return out


def merge_segments(segments: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Aggregate fine segments into a small fixed set of task stages.

    Fine-grained traces often alternate between locate/diagnose/fix many times.
    For SFT subproblem construction we want coarse, purposeful stages rather
    than exact chronological micro-steps, so this groups by task category and
    preserves the first-to-last turn span for each category.
    """

    by_category: dict[str, dict[str, Any]] = {}

    def start(seg: dict[str, Any], name: str) -> dict[str, Any]:
        name = COARSE_MAP.get(seg["phase"], "diagnose")
        return {
            "category": name,
            "goal": COARSE_GOALS[name],
            "turn_range": list(seg["turn_range"]),
            "source_phases": [seg["phase"]],
            "commands": list(seg.get("commands", [])),
            "tool_names": list(seg.get("tool_names", [])),
            "files": list(seg.get("files", [])),
            "assistant_intents": list(seg.get("assistant_intents", [])),
            "observations": list(seg.get("observations", [])),
        }

    for seg in segments:
        name = COARSE_MAP.get(seg["phase"], "diagnose")
        current = by_category.get(name)
        if current is None:
            by_category[name] = start(seg, name)
            continue

        current["turn_range"][0] = min(current["turn_range"][0], seg["turn_range"][0])
        current["turn_range"][1] = max(current["turn_range"][1], seg["turn_range"][1])
        current["source_phases"].append(seg["phase"])
        current["commands"].extend(seg.get("commands", []))
        current["tool_names"].extend(seg.get("tool_names", []))
        current["files"].extend(seg.get("files", []))
        current["assistant_intents"].extend(seg.get("assistant_intents", []))
        current["observations"].extend(seg.get("observations", []))

    ordered = []
    for name in ["understand", "locate", "diagnose", "fix", "verify", "finalize"]:
        if name in by_category:
            ordered.append(finish(by_category[name]))
    return ordered


def finish(seg: dict[str, Any]) -> dict[str, Any]:
    seg["source_phases"] = dedupe(seg["source_phases"])
    seg["commands"] = dedupe(seg["commands"], 12)
    seg["tool_names"] = dedupe(seg["tool_names"], 12)
    seg["files"] = dedupe(seg["files"], 12)
    seg["assistant_intents"] = dedupe(seg["assistant_intents"], 6)
    seg["observations"] = dedupe(seg["observations"], 6)
    return seg


def build_training_views(item: dict[str, Any], coarse_segments: list[dict[str, Any]]) -> dict[str, Any]:
    qa = []
    question_trajectory = []
    for seg in coarse_segments:
        qa.append(
            {
                "question": f"In repository {item['repo']}, how should an agent {seg['goal'][0].lower() + seg['goal'][1:]}",
                "answer": {
                    "category": seg["category"],
                    "files": seg["files"],
                    "commands": seg["commands"],
                    "evidence": seg["observations"],
                    "next_action_hint": next_action_hint(seg["category"]),
                },
            }
        )
        question_trajectory.append(
            {
                "question": f"For task {item['instance_id']} in {item['repo']}, perform the {seg['category']} stage.",
                "category": seg["category"],
                "turn_range": seg["turn_range"],
                "source_phases": seg["source_phases"],
                "trajectory_outline": {
                    "intents": seg["assistant_intents"],
                    "tools": seg["tool_names"],
                    "commands": seg["commands"],
                    "observations": seg["observations"],
                },
            }
        )
    return {"qa": qa, "question_trajectory": question_trajectory}


def next_action_hint(category: str) -> str:
    return {
        "understand": "Start repository exploration.",
        "locate": "Inspect the most relevant code paths.",
        "diagnose": "Decide the minimal edit needed to fix the root cause.",
        "fix": "Run targeted verification.",
        "verify": "Review diff and clean temporary files.",
        "finalize": "Stop or submit the final solution.",
    }[category]


def main() -> int:
    items = json.loads(IN_FILE.read_text())
    output = []
    category_counts = Counter()
    before = 0
    after = 0
    for item in items:
        coarse_segments = merge_segments(item["segments"])
        before += len(item["segments"])
        after += len(coarse_segments)
        category_counts.update(seg["category"] for seg in coarse_segments)
        converted = {
            **{k: item[k] for k in [
                "config",
                "scaffold",
                "model_family",
                "thinking_mode",
                "instance_id",
                "repo",
                "language",
                "license",
                "trajectory_id",
                "resolved",
                "trajectory_len",
                "task",
                "model_patch_head",
            ]},
            "coarse_segments": coarse_segments,
            "training_views": build_training_views(item, coarse_segments),
        }
        output.append(converted)

    report = {
        "items": len(items),
        "fine_segments": before,
        "coarse_segments": after,
        "avg_fine_segments_per_trace": before / len(items),
        "avg_coarse_segments_per_trace": after / len(items),
        "category_counts": dict(category_counts),
        "categories": COARSE_GOALS,
    }
    OUT_FILE.write_text(json.dumps(output, indent=2, ensure_ascii=False))
    REPORT_FILE.write_text(json.dumps(report, indent=2, ensure_ascii=False))
    print(json.dumps(report, indent=2, ensure_ascii=False))
    return 0


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