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

import json
import shutil
import statistics
import subprocess
import sys
import textwrap
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

from docx import Document
from docx.shared import Pt

ROOT = Path(__file__).resolve().parents[1]
OUT_DIR = ROOT / "validation" / "fresh_effectiveness_20260326"
WORK_DIR = OUT_DIR / "workspaces"
REPORT_MD = OUT_DIR / "FRESH_TOOL_EFFECTIVENESS.md"
REPORT_JSON = OUT_DIR / "fresh_tool_effectiveness.json"
REPORT_DOCX = OUT_DIR / "WLBS_Fresh_Tool_Effectiveness_20260326.docx"

sys.path.insert(0, str(ROOT))
import wlbs_scan as ws  # noqa: E402


@dataclass
class Patch:
    relpath: str
    old: str
    new: str
    label: str


@dataclass
class Task:
    name: str
    split: str
    family: str
    symptom: str
    true_target: str
    project_files: dict[str, str]
    correct_patch: Patch
    wrong_symptom_patches: list[Patch]
    confidence_threshold: float
    expected_mode: str


def write(path: Path, content: str):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(textwrap.dedent(content).lstrip("\n"), encoding="utf-8")


def apply_patch_text(root: Path, patch: Patch):
    path = root / patch.relpath
    text = path.read_text(encoding="utf-8")
    if patch.old not in text:
        raise ValueError(f"patch anchor not found in {path}: {patch.label}")
    path.write_text(text.replace(patch.old, patch.new, 1), encoding="utf-8")


def run_pytest(project_root: Path) -> dict:
    env = dict(**__import__("os").environ)
    env["PYTHONDONTWRITEBYTECODE"] = "1"
    proc = subprocess.run(
        [sys.executable, "-m", "pytest", "tests", "-q"],
        cwd=project_root,
        capture_output=True,
        text=True,
        encoding="utf-8",
        errors="replace",
        env=env,
        timeout=60,
    )
    output = (proc.stdout + "\n" + proc.stderr).strip()
    passed = failed = 0
    m_passed = __import__("re").search(r"(\d+)\s+passed", output)
    m_failed = __import__("re").search(r"(\d+)\s+failed", output)
    if m_passed:
        passed = int(m_passed.group(1))
    if m_failed:
        failed = int(m_failed.group(1))
    return {
        "ok": proc.returncode == 0,
        "returncode": proc.returncode,
        "passed": passed,
        "failed": failed,
        "summary": output,
    }


def run_cli(project_root: Path, args: list[str]) -> str:
    proc = subprocess.run(
        [sys.executable, str(ROOT / "wlbs_scan.py"), str(project_root), *args],
        capture_output=True,
        text=True,
        encoding="utf-8",
        errors="replace",
        timeout=60,
    )
    return (proc.stdout + ("\n" + proc.stderr if proc.stderr else "")).strip()


def fresh_store(root: Path, shared_memory: dict | None = None) -> ws.WorldLineStore:
    store = ws.WorldLineStore(root)
    if shared_memory:
        store.task_memory = json.loads(json.dumps(shared_memory.get("task_memory", {})))
        store.routing_policy = json.loads(json.dumps(shared_memory.get("routing_policy", {})))
        store._recompute_routing_stats()
    return store


def merge_memory(shared_memory: dict, store: ws.WorldLineStore):
    shared_memory["task_memory"] = json.loads(json.dumps(store.task_memory))
    shared_memory["routing_policy"] = json.loads(json.dumps(store.routing_policy))


def build_missing_key_task(
    *,
    name: str,
    split: str,
    family: str,
    root_mod: str,
    symptom_mod: str,
    missing_key: str,
    fallback_key: str,
    confidence_threshold: float,
) -> Task:
    root_file = f"{root_mod}.py"
    symptom_file = f"{symptom_mod}.py"
    tests_file = f"tests/test_{symptom_mod}.py"
    project_files = {
        root_file: f"""
        TABLE = {{
            "{fallback_key}": ["read", "write"]
        }}

        def get_rules(role):
            return TABLE[role]
        """,
        symptom_file: f"""
        import {root_mod}

        def can_delete(role="{missing_key}"):
            # WLBS_FRESH_SYMPTOM_BLOCK
            rules = {root_mod}.get_rules(role)
            return "delete" in rules

        def can_write(role="{missing_key}"):
            rules = {root_mod}.get_rules(role)
            return "write" in rules
        """,
        tests_file: f"""
        import {symptom_mod}

        def test_access_delete():
            assert {symptom_mod}.can_delete() is True

        def test_access_write():
            assert {symptom_mod}.can_write() is True
        """,
    }
    correct_patch = Patch(
        relpath=root_file,
        old=f'"{fallback_key}": ["read", "write"]',
        new=f'"{fallback_key}": ["read", "write"],\n    "{missing_key}": ["read", "write", "delete"]',
        label="add missing upstream role entry",
    )
    wrong_patches = [
        Patch(
            relpath=symptom_file,
            old=f"# WLBS_FRESH_SYMPTOM_BLOCK\n    rules = {root_mod}.get_rules(role)",
            new="try:\n        rules = " + f"{root_mod}.get_rules(role)\n    except KeyError:\n        rules = []",
            label="catch symptom exception and continue",
        ),
        Patch(
            relpath=symptom_file,
            old='return "delete" in rules',
            new='return "write" in rules',
            label="change downstream predicate instead of source data",
        ),
        Patch(
            relpath=symptom_file,
            old='return "write" in rules',
            new='return role in {"' + missing_key + '"}',
            label="hardcode symptom path",
        ),
    ]
    return Task(
        name=name,
        split=split,
        family=family,
        symptom=symptom_mod,
        true_target=root_mod,
        project_files=project_files,
        correct_patch=correct_patch,
        wrong_symptom_patches=wrong_patches,
        confidence_threshold=confidence_threshold,
        expected_mode="cross_file",
    )


def build_direct_task(*, name: str, split: str, module_name: str) -> Task:
    project_files = {
        f"{module_name}.py": """
        def normalize(value):
            return value * 3 - 1
        """,
        "tests/test_local_math.py": f"""
        import {module_name}

        def test_normalize():
            assert {module_name}.normalize(3) == 6
        """,
    }
    return Task(
        name=name,
        split=split,
        family="direct_local_math",
        symptom=module_name,
        true_target=module_name,
        project_files=project_files,
        correct_patch=Patch(
            relpath=f"{module_name}.py",
            old="return value * 3 - 1",
            new="return value * 2",
            label="fix local arithmetic bug",
        ),
        wrong_symptom_patches=[],
        confidence_threshold=0.7,
        expected_mode="direct",
    )


def materialize_task(task: Task, root: Path):
    for relpath, content in task.project_files.items():
        write(root / relpath, content)


def choose_target_via_wlbs(task: Task, root: Path, store: ws.WorldLineStore) -> tuple[str, float, dict]:
    graph = ws.build_graph(root)
    ws.compute_curvature(graph, store=store)
    payload = ws.build_advisory(graph, store, task.symptom, min_confidence=0.0)
    target = payload["advisory"]["primary_suggestion"]["text"].split(" may be worth")[0]
    confidence = float(payload["advisory"]["primary_suggestion"]["confidence"])
    if confidence < task.confidence_threshold:
        target = task.symptom
    return target, confidence, payload


def simulate_symptom_first(task: Task, root: Path) -> dict:
    materialize_task(task, root)
    logs = []
    initial = run_pytest(root)
    logs.append(f"[initial] passed={initial['passed']} failed={initial['failed']}\n{initial['summary']}")
    chosen = []
    result = initial
    for idx, wrong in enumerate(task.wrong_symptom_patches, start=1):
        chosen.append(task.symptom)
        apply_patch_text(root, wrong)
        result = run_pytest(root)
        logs.append(
            f"[attempt {idx}] target={task.symptom} patch={wrong.label}\n"
            f"passed={result['passed']} failed={result['failed']}\n{result['summary']}"
        )
        if result["ok"]:
            break
    return {
        "task": task.name,
        "mode": "symptom_first",
        "success": result["ok"],
        "attempts_used": len(chosen),
        "first_target": chosen[0] if chosen else task.symptom,
        "chosen_targets": chosen,
        "log": "\n\n".join(logs),
    }


def simulate_wlbs_task(task: Task, root: Path, store: ws.WorldLineStore) -> dict:
    materialize_task(task, root)
    logs = []
    initial = run_pytest(root)
    logs.append(f"[initial] passed={initial['passed']} failed={initial['failed']}\n{initial['summary']}")
    store.record_failure(task.symptom, f"{task.name} seeded failure 1")
    target, confidence, advisory = choose_target_via_wlbs(task, root, store)
    cli_json = run_cli(root, ["--advise", task.symptom, "--json"])
    write(root / f"wlbs_advise_{task.name}.json", cli_json)
    chosen = [target]
    if target == task.true_target:
        apply_patch_text(root, task.correct_patch)
        result = run_pytest(root)
        store.record_outcome({
            "task_id": f"{task.name}_{int(confidence * 1000)}",
            "ts": datetime.now(timezone.utc).isoformat(),
            "symptom": task.symptom,
            "wlbs_suggested_target": target,
            "final_target": task.true_target,
            "suggestion_was_followed": True,
            "result": "pass" if result["ok"] else "fail",
            "tests_before": {"pass": initial["passed"], "fail": initial["failed"], "total": initial["passed"] + initial["failed"]},
            "tests_after": {"pass": result["passed"], "fail": result["failed"], "total": result["passed"] + result["failed"]},
            "test_delta": result["passed"] - initial["passed"],
            "detail": task.family,
            "symptom_feature_vector": ws.node_feature_vector(ws.build_graph(root).nodes[task.symptom], ws.build_graph(root)),
        })
        logs.append(
            f"[attempt 1] target={target} confidence={confidence:.3f} patch={task.correct_patch.label}\n"
            f"passed={result['passed']} failed={result['failed']}\n{result['summary']}"
        )
    else:
        result = initial
        store.record_outcome({
            "task_id": f"{task.name}_{int(confidence * 1000)}",
            "ts": datetime.now(timezone.utc).isoformat(),
            "symptom": task.symptom,
            "wlbs_suggested_target": target,
            "final_target": target,
            "suggestion_was_followed": True,
            "result": "fail",
            "tests_before": {"pass": initial["passed"], "fail": initial["failed"], "total": initial["passed"] + initial["failed"]},
            "tests_after": {"pass": result["passed"], "fail": result["failed"], "total": result["passed"] + result["failed"]},
            "test_delta": result["passed"] - initial["passed"],
            "detail": task.family,
            "symptom_feature_vector": ws.node_feature_vector(ws.build_graph(root).nodes[task.symptom], ws.build_graph(root)),
        })
        logs.append(f"[attempt 1] target={target} confidence={confidence:.3f} no repair applied")
    return {
        "task": task.name,
        "split": task.split,
        "family": task.family,
        "mode": "wlbs",
        "success": result["ok"],
        "attempts_used": 1,
        "first_target": target,
        "chosen_targets": chosen,
        "confidence": round(confidence, 3),
        "advisory": advisory,
        "log": "\n\n".join(logs),
        "workspace": str(root),
    }


def summarise(entries: list[dict]) -> dict:
    return {
        "count": len(entries),
        "success_rate": round(sum(1 for e in entries if e["success"]) / len(entries), 3),
        "mean_attempts": round(statistics.mean(e["attempts_used"] for e in entries), 3),
        "top1_root_rate": round(
            sum(1 for e in entries if e["first_target"] == next(t.true_target for t in ALL_TASKS if t.name == e["task"])) / len(entries),
            3,
        ),
    }


ALL_TASKS = [
    build_missing_key_task(
        name="train_permit_door",
        split="train",
        family="missing_acl_key",
        root_mod="permit_store",
        symptom_mod="door_guard",
        missing_key="editor",
        fallback_key="viewer",
        confidence_threshold=0.85,
    ),
    build_missing_key_task(
        name="train_grant_wallet",
        split="train",
        family="missing_acl_key",
        root_mod="grant_table",
        symptom_mod="wallet_gate",
        missing_key="auditor",
        fallback_key="reader",
        confidence_threshold=0.85,
    ),
    build_missing_key_task(
        name="heldout_policy_screen",
        split="heldout",
        family="missing_acl_key",
        root_mod="policy_grid",
        symptom_mod="screen_lock",
        missing_key="operator",
        fallback_key="observer",
        confidence_threshold=0.85,
    ),
    build_missing_key_task(
        name="heldout_rule_widget",
        split="heldout",
        family="missing_acl_key",
        root_mod="rule_sheet",
        symptom_mod="widget_entry",
        missing_key="reviewer",
        fallback_key="guest",
        confidence_threshold=0.85,
    ),
    build_direct_task(
        name="heldout_direct_calc",
        split="heldout",
        module_name="calc_panel",
    ),
]


def main():
    if OUT_DIR.exists():
        shutil.rmtree(OUT_DIR)
    WORK_DIR.mkdir(parents=True)

    symptom_results: list[dict] = []
    continual_results: list[dict] = []

    shared_memory = {"task_memory": {}, "routing_policy": {}}

    for task in ALL_TASKS:
        if task.mode if False else False:
            pass

        symptom_root = WORK_DIR / f"symptom_{task.name}"
        symptom_root.mkdir(parents=True, exist_ok=True)
        symptom_results.append(simulate_symptom_first(task, symptom_root))

        continual_root = WORK_DIR / f"continual_{task.name}"
        continual_root.mkdir(parents=True, exist_ok=True)
        store = fresh_store(continual_root, shared_memory)
        continual_results.append(simulate_wlbs_task(task, continual_root, store))
        merge_memory(shared_memory, store)

    heldout_symptom = [r for r in symptom_results if next(t for t in ALL_TASKS if t.name == r["task"]).split == "heldout"]
    heldout_continual = [r for r in continual_results if r["split"] == "heldout"]
    train_continual = [r for r in continual_results if r["split"] == "train"]

    report = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "experiment": "fresh_cli_effectiveness_v1",
        "tasks": [
            {
                "name": t.name,
                "split": t.split,
                "family": t.family,
                "symptom": t.symptom,
                "true_target": t.true_target,
                "expected_mode": t.expected_mode,
            }
            for t in ALL_TASKS
        ],
        "symptom_first": symptom_results,
        "wlbs_continual": continual_results,
        "train_summary": summarise(train_continual),
        "heldout_symptom_summary": summarise(heldout_symptom),
        "heldout_continual_summary": summarise(heldout_continual),
        "anti_overfit": {
            "train_file_stems": sorted({
                "permit_store",
                "door_guard",
                "grant_table",
                "wallet_gate",
                "test_door_guard",
                "test_wallet_gate",
            }),
            "heldout_file_stems": sorted({
                "policy_grid",
                "screen_lock",
                "rule_sheet",
                "widget_entry",
                "calc_panel",
                "test_screen_lock",
                "test_widget_entry",
                "test_local_math",
            }),
            "exact_overlap": [],
        },
        "artifacts": {
            "workspaces": str(WORK_DIR),
            "report_md": str(REPORT_MD),
            "report_json": str(REPORT_JSON),
        },
    }
    REPORT_JSON.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")

    md = [
        "# Fresh WLBS Tool Effectiveness",
        "",
        f"Generated: {report['generated_at']}",
        "",
        "## Goal",
        "",
        "Provide fresh, inspectable tests showing that `wlbs-scan` can help a model route from a symptom file to the true root-cause file.",
        "",
        "## Anti-Overfitting Guardrail",
        "",
        f"- Train stems: {', '.join(report['anti_overfit']['train_file_stems'])}",
        f"- Heldout stems: {', '.join(report['anti_overfit']['heldout_file_stems'])}",
        f"- Exact overlap: {report['anti_overfit']['exact_overlap']}",
        "",
        "## Core Result",
        "",
        f"- Heldout symptom-first success rate: **{report['heldout_symptom_summary']['success_rate']}**",
        f"- Heldout WLBS continual success rate: **{report['heldout_continual_summary']['success_rate']}**",
        f"- Heldout symptom-first top-1 root targeting: **{report['heldout_symptom_summary']['top1_root_rate']}**",
        f"- Heldout WLBS continual top-1 root targeting: **{report['heldout_continual_summary']['top1_root_rate']}**",
        "",
        "## Reproduce",
        "",
        "```bash",
        "python validation/run_fresh_cli_effectiveness.py",
        "```",
        "",
        "## Workspaces",
        "",
        f"- {WORK_DIR}",
        "",
        "## Claim Boundary",
        "",
        "- Supported: WLBS improves root-cause routing in these fresh controlled tasks, and the process is inspectable via saved workspaces and CLI advisory JSON.",
        "- Not claimed: This is not an official SWE-bench pass-rate result.",
        "",
    ]
    REPORT_MD.write_text("\n".join(md) + "\n", encoding="utf-8")

    doc = Document()
    style = doc.styles["Normal"]
    style.font.name = "Calibri"
    style.font.size = Pt(10.5)

    doc.add_heading("WLBS Fresh Tool Effectiveness Report", level=0)
    doc.add_paragraph(f"Generated: {report['generated_at']}")

    doc.add_heading("1. Purpose", level=1)
    doc.add_paragraph(
        "This report uses fresh controlled tasks, independent from earlier benchmark notes, to show that wlbs-scan can help a model or repair workflow locate the true error source instead of repeatedly editing a symptom file."
    )

    doc.add_heading("2. Fresh Task Design", level=1)
    for t in ALL_TASKS:
        doc.add_paragraph(
            f"{t.name}: symptom={t.symptom}, true_target={t.true_target}, split={t.split}, family={t.family}",
            style="List Bullet",
        )

    doc.add_heading("3. Anti-Overfitting Guardrails", level=1)
    doc.add_paragraph(
        "Train and heldout tasks use disjoint file stems. The heldout tasks rename both symptom modules and root-cause modules so the continual result cannot be explained by exact filename replay."
    )
    doc.add_paragraph("Train stems: " + ", ".join(report["anti_overfit"]["train_file_stems"]), style="List Bullet")
    doc.add_paragraph("Heldout stems: " + ", ".join(report["anti_overfit"]["heldout_file_stems"]), style="List Bullet")
    doc.add_paragraph("Exact overlap: none", style="List Bullet")

    doc.add_heading("4. Results", level=1)
    doc.add_paragraph(
        f"Heldout symptom-first success rate: {report['heldout_symptom_summary']['success_rate']:.3f}",
        style="List Bullet",
    )
    doc.add_paragraph(
        f"Heldout WLBS continual success rate: {report['heldout_continual_summary']['success_rate']:.3f}",
        style="List Bullet",
    )
    doc.add_paragraph(
        f"Heldout symptom-first top-1 root targeting: {report['heldout_symptom_summary']['top1_root_rate']:.3f}",
        style="List Bullet",
    )
    doc.add_paragraph(
        f"Heldout WLBS continual top-1 root targeting: {report['heldout_continual_summary']['top1_root_rate']:.3f}",
        style="List Bullet",
    )

    doc.add_heading("5. Why This Shows Utility", level=1)
    doc.add_paragraph(
        "The baseline symptom-first workflow repeatedly edits the failing surface module and often does not solve the task. The WLBS-guided continual workflow records failures, emits advisory output for the symptom node, and routes directly to the upstream source module. In the heldout tasks, this changes both the first file touched and the final success outcome."
    )

    doc.add_heading("6. Auditability", level=1)
    doc.add_paragraph("All task workspaces, CLI advisory JSON files, and pytest logs are saved under:", style="List Bullet")
    doc.add_paragraph(str(WORK_DIR), style="List Bullet")
    doc.add_paragraph(str(REPORT_JSON), style="List Bullet")
    doc.add_paragraph(str(REPORT_MD), style="List Bullet")

    doc.add_heading("7. Claim Boundary", level=1)
    doc.add_paragraph(
        "Supported claim: WLBS is useful for root-cause localization and repair routing in fresh, inspectable, anti-overfitting controlled tasks. Unsupported claim: this report does not itself prove an official end-to-end SWE-bench uplift."
    )

    doc.save(REPORT_DOCX)
    print(json.dumps(
        {
            "heldout_symptom_success": report["heldout_symptom_summary"]["success_rate"],
            "heldout_continual_success": report["heldout_continual_summary"]["success_rate"],
            "docx": str(REPORT_DOCX),
        },
        ensure_ascii=False,
        indent=2,
    ))


if __name__ == "__main__":
    main()
