#!/usr/bin/env python3
"""Triage rejected samples into variation repair decisions.

The goal is to use rejected_sample_10k as evidence for fixing chart variations
instead of post-processing individual images. The script joins each rejected
sample with:

- its rejection reasons and blockers,
- one GPT-improved reference image, and
- the source chart variation/template path when available.

It then emits variation-level repair briefs and icon/image generation batch
prompts. The classifier is intentionally conservative and rule-based so that it
is auditable before LLM-assisted code editing is introduced.
"""

from __future__ import annotations

import argparse
import json
import re
from collections import Counter, defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable


PROJECT_ROOT = Path(__file__).resolve().parents[1]
WORKSPACE_ROOT = PROJECT_ROOT.parent
DEFAULT_REJECTED_ROOT = WORKSPACE_ROOT / "rejected_sample_10k"
GPT_SUFFIX_RE = re.compile(r"^(?P<base>.+)_gpt_image_2_(?P<variant>\d+)\.(?:png|jpg|jpeg|webp)$", re.I)


RULES = {
    "asset_semantics": [
        "icon",
        "illustration",
        "semantically unrelated",
        "topic mismatch",
        "unrelated to the topic",
        "generic",
        "decorative",
        "cartoon",
    ],
    "asset_obstruction": [
        "obstruct",
        "sits directly",
        "overlap",
        "cover",
        "blocks",
        "visual clutter",
    ],
    "ordering": [
        "chronological",
        "not in order",
        "out of order",
        "jumps between",
        "sorted",
        "sequence",
    ],
    "text_readability": [
        "text",
        "label",
        "readability",
        "unreadable",
        "overlapping",
        "cluttered",
        "rotated",
        "upside-down",
        "sideways",
        "cropped",
        "small",
    ],
    "axis_scale": [
        "axis",
        "scale",
        "tick",
        "repeated",
        "range",
        "number '20'",
    ],
    "chart_semantics": [
        "chart type",
        "radial/polar",
        "polar chart",
        "does not effectively",
        "inefficient",
        "confusing",
        "meaningless",
        "unusable",
        "hard to follow",
    ],
    "layout_coherence": [
        "layout",
        "collage",
        "disconnected",
        "cohesive",
        "hierarchy",
        "space",
        "crowded",
    ],
}


@dataclass
class SampleTriage:
    sample_id: str
    chart_variation: str
    chart_type: str
    decision: str
    categories: list[str]
    blockers: list[str]
    reasons: list[str]
    original_image: str | None
    gpt_reference: str | None
    scores: dict[str, Any]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Build repair briefs from rejected ChartGalaxy samples.")
    parser.add_argument("--rejected-root", default=str(DEFAULT_REJECTED_ROOT))
    parser.add_argument(
        "--output-dir",
        default=None,
        help="Default: <rejected-root>/variation_repair_triage",
    )
    parser.add_argument(
        "--gpt-variant",
        type=int,
        default=0,
        help="Use exactly one GPT reference variant per sample.",
    )
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--min-variation-count", type=int, default=1)
    parser.add_argument(
        "--top",
        type=int,
        default=80,
        help="Max number of variations in the markdown report.",
    )
    return parser.parse_args()


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


def write_json(path: Path, data: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        json.dump(data, handle, indent=2, ensure_ascii=False)
        handle.write("\n")


def write_jsonl(path: Path, records: Iterable[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for record in records:
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")


def read_manifest_refs(rejected_root: Path, variant: int) -> dict[str, str]:
    refs: dict[str, str] = {}
    for manifest_path in [
        rejected_root / "gpt_image_2_improved" / "manifest.jsonl",
        rejected_root / "gpt_image_2_improved" / "manifest.jsonl.500w",
    ]:
        if not manifest_path.exists():
            continue
        with manifest_path.open("r", encoding="utf-8") as handle:
            for line in handle:
                line = line.strip()
                if not line:
                    continue
                try:
                    record = json.loads(line)
                except json.JSONDecodeError:
                    continue
                input_path = Path(record.get("input", ""))
                sample_id = input_path.stem
                outputs = record.get("outputs") or []
                chosen = None
                for output in outputs:
                    match = GPT_SUFFIX_RE.match(Path(output).name)
                    if match and int(match.group("variant")) == variant:
                        chosen = output
                        break
                if chosen:
                    chosen_path = Path(chosen)
                    candidates = []
                    if chosen_path.is_absolute():
                        candidates.append(chosen_path)
                    else:
                        candidates.append((WORKSPACE_ROOT / chosen_path).resolve())
                        candidates.append((rejected_root.parent / chosen_path).resolve())
                    candidates.append(rejected_root / "gpt_image_2_improved" / Path(chosen).name)
                    for candidate in candidates:
                        if candidate.exists():
                            refs[sample_id] = str(candidate.resolve())
                            break
    return refs


def find_gpt_ref_by_scan(rejected_root: Path, sample_id: str, variant: int) -> str | None:
    suffix = f"_gpt_image_2_{variant:02d}.png"
    candidates = sorted((rejected_root / "gpt_image_2_improved").glob(f"{sample_id}{suffix}"))
    if candidates:
        return str(candidates[0].resolve())
    return None


def template_paths_by_variation() -> dict[str, list[str]]:
    roots = [
        PROJECT_ROOT / "modules" / "chart_engine" / "template",
    ]
    mapping: dict[str, list[str]] = defaultdict(list)
    for root in roots:
        if not root.exists():
            continue
        for path in root.rglob("*"):
            if path.suffix.lower() not in {".js", ".py"}:
                continue
            if path.name.startswith("__"):
                continue
            mapping[path.stem].append(str(path.resolve()))
    return dict(mapping)


def match_categories(reasons: list[str], blockers: list[str]) -> list[str]:
    text = " ".join(reasons + blockers).lower()
    categories: list[str] = []
    for category, keywords in RULES.items():
        if any(keyword in text for keyword in keywords):
            categories.append(category)
    if not categories:
        categories.append("needs_manual_review")
    return categories


def decide(categories: list[str], blockers: list[str]) -> str:
    category_set = set(categories)
    blocker_set = set(blockers)

    asset_only = category_set <= {"asset_semantics", "asset_obstruction", "layout_coherence"}
    if asset_only and ("topic_relevance" in blocker_set or "icon_illustration_quality" in blocker_set):
        return "repair_assets"

    if category_set & {"ordering", "text_readability", "axis_scale"}:
        if "chart_semantics" in category_set and "overall_quality" in blocker_set:
            return "repair_or_restrict_variation"
        return "repair_template"

    if "asset_semantics" in category_set or "asset_obstruction" in category_set:
        return "repair_assets_and_layout"

    if "chart_semantics" in category_set:
        return "restrict_or_abandon_variation"

    if "layout_coherence" in category_set:
        return "repair_layout"

    return "manual_review"


def topic_from_sample_id(sample_id: str) -> str:
    if "__" not in sample_id:
        return ""
    tail = sample_id.split("__", 1)[1]
    tail = re.sub(r"_(scenario|temporal|categorical|numerical|tp|xy|xyear|y|group).*", "", tail)
    tail = re.sub(r"_\d+$", "", tail)
    return tail.replace("_", " ").strip()


def action_brief_for_categories(categories: Iterable[str], decision: str) -> list[str]:
    cats = set(categories)
    actions: list[str] = []
    if "ordering" in cats:
        actions.append("Sort temporal/category inputs before rendering; never let radial/axis labels follow raw unsorted order.")
    if "text_readability" in cats:
        actions.append("Reduce rotated labels, add collision checks, and fall back to outside labels/legends when label density is high.")
    if "axis_scale" in cats:
        actions.append("Fix tick generation/formatting so axis labels are unique, monotonic, and data-range aware.")
    if "asset_semantics" in cats:
        actions.append("Replace retrieved one-off icons with domain/topic-aware generated icon sets.")
    if "asset_obstruction" in cats:
        actions.append("Keep illustration/image masks outside the chart data region; enforce no-overlap constraints.")
    if "layout_coherence" in cats:
        actions.append("Align supporting panels with the chart and remove disconnected side boxes when they do not encode data.")
    if "chart_semantics" in cats:
        actions.append("Keep the chart type unchanged; if the defect requires a different visual encoding, restrict or abandon this variation for that data shape.")
    if not actions:
        actions.append(f"Manual review required before modifying the variation ({decision}).")
    return actions


def make_icon_job(variation: str, samples: list[SampleTriage], template_paths: list[str]) -> dict[str, Any] | None:
    asset_samples = [
        sample
        for sample in samples
        if sample.decision in {"repair_assets", "repair_assets_and_layout"}
        or "asset_semantics" in sample.categories
        or "asset_obstruction" in sample.categories
    ]
    if not asset_samples:
        return None

    topics = []
    for sample in asset_samples:
        topic = topic_from_sample_id(sample.sample_id)
        if topic and topic not in topics:
            topics.append(topic)
        if len(topics) >= 24:
            break

    prompt_topics = topics or [variation.replace("_", " ")]
    prompt = (
        "Create one coherent infographic icon sheet in a single consistent visual system. "
        "Use flat vector-like pictograms with matching stroke weight, palette, lighting, and perspective. "
        "No text, letters, numbers, labels, watermarks, or captions. "
        "Use a pure white background. Arrange icons in a clean grid, one icon per cell. "
        "Topics: "
        + "; ".join(prompt_topics)
        + "."
    )
    return {
        "chart_variation": variation,
        "template_paths": template_paths,
        "sample_count": len(asset_samples),
        "topics": prompt_topics,
        "prompt": prompt,
        "reference_images": [sample.gpt_reference for sample in asset_samples[:3] if sample.gpt_reference],
        "original_images": [sample.original_image for sample in asset_samples[:3] if sample.original_image],
    }


def make_llm_review_job(summary: dict[str, Any]) -> dict[str, Any]:
    images = []
    for sample in summary["representative_samples"][:3]:
        if sample.get("original_image"):
            images.append({"role": "original_rejected", "path": sample["original_image"], "sample_id": sample["sample_id"]})
        if sample.get("gpt_reference"):
            images.append({"role": "gpt_style_reference", "path": sample["gpt_reference"], "sample_id": sample["sample_id"]})

    prompt = f"""You are repairing a ChartGalaxy chart variation, not post-processing a single output.

Variation: {summary['chart_variation']}
Chart type: {summary['chart_type']}
Template paths: {summary['template_paths']}
Current triage decision: {summary['dominant_decision']}
Aggregated blockers: {summary['blocker_counts']}
Aggregated defect categories: {summary['category_counts']}

Use the original rejected images to identify actual failures. Use GPT-improved images only as style references for layout, visual hierarchy, coherent decoration, and icon style. Do not copy GPT-rendered chart data, numeric values, labels, or geometry.

Do not change the chart type or visual encoding of the variation. A line chart must remain a line chart, a gauge must remain the same gauge family, a proportional-area variation must remain the same proportional-area mark family, etc. If the only good fix would be to turn this into another chart type, return restrict_variation or abandon_variation and describe which data shapes should be routed elsewhere.

Return JSON only:
{{
  "decision": "repair_template | repair_assets | restrict_variation | abandon_variation | manual_review",
  "fixability_reason": "...",
  "template_patch_brief": ["concrete code-level changes to the variation/template"],
  "data_compatibility_rules": ["rules for when this variation should not be selected"],
  "asset_generation_brief": {{
    "needed": true,
    "icon_sheet_prompt": "single coherent image-generation prompt if icons/images should be regenerated as a set",
    "placement_rules": ["where images/icons may be placed relative to chart data"]
  }},
  "validation_checks": ["tests or visual checks that must pass after patching"]
}}
"""
    return {
        "chart_variation": summary["chart_variation"],
        "chart_type": summary["chart_type"],
        "template_paths": summary["template_paths"],
        "dominant_decision": summary["dominant_decision"],
        "representative_samples": summary["representative_samples"][:3],
        "images": images,
        "prompt": prompt,
    }


def summarize_variation(
    variation: str,
    samples: list[SampleTriage],
    template_paths: list[str],
) -> dict[str, Any]:
    decision_counts = Counter(sample.decision for sample in samples)
    category_counts = Counter(category for sample in samples for category in sample.categories)
    blocker_counts = Counter(blocker for sample in samples for blocker in sample.blockers)
    dominant_decision = decision_counts.most_common(1)[0][0]
    if "repair_template" in decision_counts or "repair_assets" in decision_counts or "repair_assets_and_layout" in decision_counts:
        dominant_decision = "repair_variation"
    elif "repair_or_restrict_variation" in decision_counts:
        dominant_decision = "repair_or_restrict_variation"
    elif "restrict_or_abandon_variation" in decision_counts:
        dominant_decision = "restrict_or_abandon_variation"

    representative = sorted(
        samples,
        key=lambda sample: (
            -len(sample.blockers),
            sample.scores.get("overall_quality", 99),
            sample.sample_id,
        ),
    )[:5]

    return {
        "chart_variation": variation,
        "chart_type": samples[0].chart_type if samples else "",
        "sample_count": len(samples),
        "dominant_decision": dominant_decision,
        "decision_counts": dict(decision_counts),
        "category_counts": dict(category_counts),
        "blocker_counts": dict(blocker_counts),
        "template_paths": template_paths,
        "repair_actions": action_brief_for_categories(category_counts.keys(), dominant_decision),
        "representative_samples": [
            {
                "sample_id": sample.sample_id,
                "decision": sample.decision,
                "categories": sample.categories,
                "blockers": sample.blockers,
                "reasons": sample.reasons,
                "original_image": sample.original_image,
                "gpt_reference": sample.gpt_reference,
            }
            for sample in representative
        ],
    }


def markdown_report(variation_summaries: list[dict[str, Any]], top: int) -> str:
    lines = [
        "# Variation Repair Triage",
        "",
        "This report groups rejected samples by chart variation and proposes whether to repair, restrict, or abandon each variation.",
        "",
    ]
    for item in variation_summaries[:top]:
        lines.append(f"## {item['chart_variation']} ({item['sample_count']} samples)")
        lines.append(f"- decision: `{item['dominant_decision']}`")
        lines.append(f"- chart_type: {item['chart_type']}")
        lines.append(f"- blockers: {item['blocker_counts']}")
        lines.append(f"- categories: {item['category_counts']}")
        if item["template_paths"]:
            lines.append("- template_paths:")
            for path in item["template_paths"]:
                lines.append(f"  - `{path}`")
        lines.append("- repair_actions:")
        for action in item["repair_actions"]:
            lines.append(f"  - {action}")
        lines.append("- representative:")
        for sample in item["representative_samples"][:3]:
            reason = sample["reasons"][0] if sample["reasons"] else ""
            lines.append(f"  - `{sample['sample_id']}`: {reason}")
            if sample.get("gpt_reference"):
                lines.append(f"    GPT ref: `{sample['gpt_reference']}`")
        lines.append("")
    return "\n".join(lines)


def main() -> int:
    args = parse_args()
    rejected_root = Path(args.rejected_root).resolve()
    output_dir = Path(args.output_dir).resolve() if args.output_dir else rejected_root / "variation_repair_triage"
    reasons_root = rejected_root / "reasons"
    images_root = rejected_root / "images"
    if not reasons_root.exists():
        raise SystemExit(f"missing reasons directory: {reasons_root}")

    gpt_refs = read_manifest_refs(rejected_root, args.gpt_variant)
    template_map = template_paths_by_variation()

    sample_records: list[SampleTriage] = []
    reason_paths = sorted(reasons_root.glob("*.json"))
    if args.limit > 0:
        reason_paths = reason_paths[: args.limit]

    for reason_path in reason_paths:
        data = load_json(reason_path)
        sample_id = data.get("sample_id") or reason_path.stem
        reasons = [str(item) for item in data.get("reject_reasons") or []]
        blockers = [str(item) for item in data.get("filter_blockers") or []]
        categories = match_categories(reasons, blockers)
        decision = decide(categories, blockers)
        original_path = images_root / f"{sample_id}.webp"
        gpt_reference = gpt_refs.get(sample_id) or find_gpt_ref_by_scan(rejected_root, sample_id, args.gpt_variant)
        sample_records.append(
            SampleTriage(
                sample_id=sample_id,
                chart_variation=str(data.get("chart_variation") or ""),
                chart_type=str(data.get("chart_type") or ""),
                decision=decision,
                categories=categories,
                blockers=blockers,
                reasons=reasons,
                original_image=str(original_path.resolve()) if original_path.exists() else None,
                gpt_reference=gpt_reference,
                scores=data.get("scores") or {},
            )
        )

    grouped: dict[str, list[SampleTriage]] = defaultdict(list)
    for sample in sample_records:
        grouped[sample.chart_variation].append(sample)

    variation_summaries = []
    for variation, samples in grouped.items():
        if len(samples) < args.min_variation_count:
            continue
        variation_summaries.append(
            summarize_variation(
                variation,
                samples,
                template_map.get(variation, []),
            )
        )
    variation_summaries.sort(
        key=lambda item: (
            -item["sample_count"],
            item["dominant_decision"],
            item["chart_variation"],
        )
    )

    icon_jobs = []
    llm_review_jobs = []
    for item in variation_summaries:
        samples = grouped[item["chart_variation"]]
        job = make_icon_job(item["chart_variation"], samples, item["template_paths"])
        if job:
            icon_jobs.append(job)
        llm_review_jobs.append(make_llm_review_job(item))

    sample_dicts = [sample.__dict__ for sample in sample_records]
    write_json(output_dir / "samples.json", sample_dicts)
    write_json(output_dir / "variation_summaries.json", variation_summaries)
    write_jsonl(output_dir / "icon_generation_jobs.jsonl", icon_jobs)
    write_jsonl(output_dir / "llm_variation_review_jobs.jsonl", llm_review_jobs)
    (output_dir / "report.md").write_text(markdown_report(variation_summaries, args.top), encoding="utf-8")
    write_json(
        output_dir / "summary.json",
        {
            "rejected_root": str(rejected_root),
            "output_dir": str(output_dir),
            "gpt_variant": args.gpt_variant,
            "sample_count": len(sample_records),
            "variation_count": len(variation_summaries),
            "icon_job_count": len(icon_jobs),
            "llm_review_job_count": len(llm_review_jobs),
            "decision_counts": dict(Counter(sample.decision for sample in sample_records)),
            "category_counts": dict(Counter(category for sample in sample_records for category in sample.categories)),
        },
    )
    print(f"samples={len(sample_records)} variations={len(variation_summaries)} icon_jobs={len(icon_jobs)}")
    print(f"report={output_dir / 'report.md'}")
    return 0


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