import os
from flask import Blueprint, jsonify
from datasets import load_dataset

bp = Blueprint("plan_revisions", __name__, url_prefix="/api/plan-revisions")

HF_REPO = "timchen0618/bcp-plan-revisions-v1"

_cache: dict | None = None


def _load():
    global _cache
    if _cache is not None:
        return _cache

    ds = load_dataset(HF_REPO, split="train")

    # Group: condition -> query_id -> sorted list of revision entries
    grouped: dict[str, dict[str, list]] = {}
    for row in ds:
        cond = row["condition"]
        qid = str(row["query_id"])
        grouped.setdefault(cond, {}).setdefault(qid, [])
        grouped[cond][qid].append({
            "revision_index": row["revision_index"],
            "source": row["source"],
            "plan_text": row["plan_text"],
            "total_revisions": row["total_revisions"],
            "status": row["status"],
        })

    # Sort revisions within each group
    for cond in grouped:
        for qid in grouped[cond]:
            grouped[cond][qid].sort(key=lambda r: r["revision_index"])

    conditions = sorted(grouped.keys())
    _cache = {"conditions": conditions, "data": grouped}
    return _cache


@bp.get("/")
def get_data():
    try:
        result = _load()
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@bp.post("/reload")
def reload_data():
    global _cache
    _cache = None
    try:
        result = _load()
        return jsonify({"status": "ok", "conditions": result["conditions"]})
    except Exception as e:
        return jsonify({"error": str(e)}), 500