#!/usr/bin/env python3
"""Derived Task-2 audit on real released VCG-Bench mxGraph diagrams.

This does not recreate the missing Task-2 release. It constructs deterministic,
machine-checkable edits from real Task-1 XML and tests execution, preservation,
and instruction discrimination for correct/no-op/wrong-target patches.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import random
import subprocess
import time
import xml.etree.ElementTree as ET
from collections import Counter
from pathlib import Path

import pandas as pd

SEED = 20260722
COLORS = {"easy": "#ff6666", "medium": "#6699ff", "hard": "#66cc99"}


def render(drawio: Path, xml_path: Path, png_path: Path) -> tuple[bool, str]:
    """Use the exact Draw.io CLI protocol used by the existing VCG audit."""
    command = [
        "xvfb-run", "-a", str(drawio), "-x", "-f", "png", "-s", "1.0",
        "-o", str(png_path), "--no-sandbox", str(xml_path),
    ]
    try:
        completed = subprocess.run(command, capture_output=True, text=True, timeout=90)
        success = completed.returncode == 0 and png_path.exists() and png_path.stat().st_size > 0
        return success, (completed.stdout + "\n" + completed.stderr).strip()[-2000:]
    except subprocess.TimeoutExpired:
        return False, "Draw.io timed out after 90 seconds"


def set_style(style: str, key: str, value: str) -> str:
    parts = [part for part in style.split(";") if part]
    mapping = {}
    order = []
    for part in parts:
        if "=" not in part:
            continue
        name, old = part.split("=", 1)
        if name not in mapping:
            order.append(name)
        mapping[name] = old
    if key not in mapping:
        order.append(key)
    mapping[key] = value
    return ";".join(f"{name}={mapping[name]}" for name in order) + ";"


def eligible_rows(frame: pd.DataFrame, per_domain: int) -> pd.DataFrame:
    rng = random.Random(SEED)
    selected = []
    for domain, group in sorted(frame.groupby("domain_l1")):
        candidates = []
        for index, row in group.iterrows():
            try:
                root = ET.fromstring(str(row.restored_xml))
            except ET.ParseError:
                continue
            vertices = [cell for cell in root.iter("mxCell") if cell.get("vertex") == "1"]
            if len(vertices) >= 2 and all(vertex.find("mxGeometry") is not None for vertex in vertices[:2]):
                candidates.append(index)
        rng.shuffle(candidates)
        if len(candidates) < per_domain:
            raise RuntimeError((domain, len(candidates)))
        selected.extend(candidates[:per_domain])
    return frame.loc[selected].reset_index(drop=True)


def apply_operations(xml: str, difficulty: str, target_position: int, marker: str) -> tuple[str, dict]:
    root = ET.fromstring(xml)
    vertices = [cell for cell in root.iter("mxCell") if cell.get("vertex") == "1"]
    target = vertices[target_position]
    geometry = target.find("mxGeometry")
    assert geometry is not None
    original = {
        "id": target.get("id"),
        "value": target.get("value", ""),
        "style": target.get("style", ""),
        "x": float(geometry.get("x", "0")),
        "y": float(geometry.get("y", "0")),
        "width": float(geometry.get("width", "0")),
        "height": float(geometry.get("height", "0")),
    }
    expected = {"target_id": target.get("id"), "color": COLORS[difficulty]}
    target.set("style", set_style(target.get("style", ""), "fillColor", COLORS[difficulty]))
    operations = ["fill_color"]
    if difficulty in {"medium", "hard"}:
        expected["value"] = f"{original['value']} [{marker}]"
        expected["width"] = round(original["width"] * 1.2, 6)
        target.set("value", expected["value"])
        geometry.set("width", f"{expected['width']:g}")
        operations.extend(["change_text", "resize_width"])
    if difficulty == "hard":
        expected["x"] = round(original["x"] + 30.0, 6)
        expected["y"] = round(original["y"] + 20.0, 6)
        expected["height"] = round(original["height"] * 1.1, 6)
        geometry.set("x", f"{expected['x']:g}")
        geometry.set("y", f"{expected['y']:g}")
        geometry.set("height", f"{expected['height']:g}")
        operations.extend(["move_x", "move_y", "resize_height"])
    expected["operations"] = operations
    expected["original_target"] = original
    return ET.tostring(root, encoding="unicode"), expected


def style_value(style: str, key: str) -> str | None:
    for part in style.split(";"):
        if part.startswith(key + "="):
            return part.split("=", 1)[1]
    return None


def exact_xdrfr(xml: str, expected: dict) -> tuple[float, list[dict]]:
    root = ET.fromstring(xml)
    target = next(cell for cell in root.iter("mxCell") if cell.get("id") == expected["target_id"])
    geometry = target.find("mxGeometry")
    assert geometry is not None
    checks = [("fill_color", style_value(target.get("style", ""), "fillColor") == expected["color"])]
    if "value" in expected:
        checks.extend(
            [
                ("change_text", target.get("value") == expected["value"]),
                ("resize_width", abs(float(geometry.get("width", "nan")) - expected["width"]) < 1e-6),
            ]
        )
    if "x" in expected:
        checks.extend(
            [
                ("move_x", abs(float(geometry.get("x", "nan")) - expected["x"]) < 1e-6),
                ("move_y", abs(float(geometry.get("y", "nan")) - expected["y"]) < 1e-6),
                ("resize_height", abs(float(geometry.get("height", "nan")) - expected["height"]) < 1e-6),
            ]
        )
    details = [{"question": name, "is_satisfied": bool(ok)} for name, ok in checks]
    return sum(ok for _, ok in checks) / len(checks), details


def cell_signatures(xml: str) -> dict[str, str]:
    root = ET.fromstring(xml)
    return {
        cell.get("id"): ET.tostring(cell, encoding="unicode")
        for cell in root.iter("mxCell")
        if cell.get("id") is not None
    }


def preservation_score(original: str, modified: str, target_id: str) -> float:
    left, right = cell_signatures(original), cell_signatures(modified)
    ids = sorted((set(left) & set(right)) - {target_id})
    return sum(left[item] == right[item] for item in ids) / len(ids) if ids else 1.0


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--parquet", type=Path, default=Path("dataset/train.parquet"))
    parser.add_argument("--drawio", type=Path, default=Path("tools/squashfs-root/drawio"))
    parser.add_argument("--output", type=Path, default=Path("results/task2_derived_audit"))
    parser.add_argument("--per-domain", type=int, default=3)
    args = parser.parse_args()
    root_dir = Path(__file__).resolve().parent
    parquet = (root_dir / args.parquet).resolve()
    drawio = (root_dir / args.drawio).resolve()
    output = (root_dir / args.output).resolve()
    xml_dir, png_dir = output / "xml", output / "rendered"
    xml_dir.mkdir(parents=True, exist_ok=True)
    png_dir.mkdir(parents=True, exist_ok=True)

    selected = eligible_rows(pd.read_parquet(parquet), args.per_domain)
    rows = []
    for sample_index, row in selected.iterrows():
        original = str(row.restored_xml)
        for difficulty in ("easy", "medium", "hard"):
            case_id = f"{row.image_id}__{difficulty}"
            correct, expected = apply_operations(original, difficulty, 0, case_id)
            wrong, _ = apply_operations(original, difficulty, 1, case_id)
            variants = {"correct_patch": correct, "no_op": original, "wrong_target": wrong}
            for variant, xml in variants.items():
                score, checks = exact_xdrfr(xml, expected)
                rows.append(
                    {
                        "case_id": case_id,
                        "image_id": str(row.image_id),
                        "domain_l1": str(row.domain_l1),
                        "domain_l2": str(row.domain_l2),
                        "difficulty": difficulty,
                        "operation_count": len(expected["operations"]),
                        "variant": variant,
                        "exact_xdrfr": score,
                        "satisfied": sum(item["is_satisfied"] for item in checks),
                        "questions": len(checks),
                        "untouched_cell_preservation": preservation_score(original, xml, expected["target_id"]),
                        "xml_parse_success": True,
                    }
                )
            xml_path = xml_dir / f"{case_id}.drawio"
            png_path = png_dir / f"{case_id}.png"
            xml_path.write_text(correct, encoding="utf-8")
            started = time.perf_counter()
            success, log = render(drawio, xml_path, png_path)
            rows[-3]["render_success"] = success
            rows[-3]["render_seconds"] = time.perf_counter() - started
            rows[-3]["render_sha256"] = sha256(png_path) if success else None
            rows[-3]["render_log_tail"] = log[-400:]
            print(f"[{sample_index + 1}/{len(selected)}] {case_id} render={success}", flush=True)

    frame = pd.DataFrame(rows)
    frame.to_csv(output / "per_task.csv", index=False)
    correct = frame[frame.variant == "correct_patch"]
    summary = {
        "scope": "derived synthetic Task-2 edits on real released Task-1 mxGraph XML; not the missing official Task-2 benchmark",
        "seed": SEED,
        "source_rows": len(selected),
        "domains": dict(Counter(selected.domain_l1)),
        "derived_edit_tasks": len(correct),
        "difficulty_counts": dict(Counter(correct.difficulty)),
        "operation_counts": sorted(correct.operation_count.unique().tolist()),
        "correct_patch": {
            "drawio_render_successes": int(correct.render_success.sum()),
            "execution_success_rate": float(correct.render_success.mean()),
            "mean_exact_xdrfr": float(correct.exact_xdrfr.mean()),
            "mean_untouched_cell_preservation": float(correct.untouched_cell_preservation.mean()),
        },
        "controls": {
            variant: {
                "mean_exact_xdrfr": float(group.exact_xdrfr.mean()),
                "perfect_instruction_following": int((group.exact_xdrfr == 1).sum()),
            }
            for variant, group in frame[frame.variant != "correct_patch"].groupby("variant")
        },
        "discrimination_gap_correct_minus_no_op": float(
            correct.exact_xdrfr.mean() - frame[frame.variant == "no_op"].exact_xdrfr.mean()
        ),
        "limitations": [
            "Task-2 source parquet and author model outputs were not released.",
            "Edits are deterministic oracle patches, not LLM-generated patches.",
            "XDRFR questions are evaluated by exact XML property checks rather than the paper's Gemini judge.",
            "Untouched-cell XML preservation is a deterministic style-preservation proxy, not VLM SCS.",
        ],
    }
    (output / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")
    assert summary["correct_patch"]["execution_success_rate"] == 1.0
    assert summary["correct_patch"]["mean_exact_xdrfr"] == 1.0
    assert summary["correct_patch"]["mean_untouched_cell_preservation"] == 1.0
    assert summary["controls"]["no_op"]["mean_exact_xdrfr"] == 0.0
    assert summary["controls"]["wrong_target"]["mean_exact_xdrfr"] == 0.0
    print(json.dumps(summary, indent=2), flush=True)


if __name__ == "__main__":
    main()