from flask import Blueprint, jsonify
from datasets import load_dataset

bp = Blueprint("traj_ext", __name__, url_prefix="/api/traj-ext")

HF_REPO = "timchen0618/bcp-traj-ext-formatted-v1"

_cache: list | None = None


def _load():
    global _cache
    if _cache is not None:
        return _cache
    ds = load_dataset(HF_REPO, split="train")
    rows = []
    for row in ds:
        rows.append({
            "query_id": str(row["query_id"]),
            "question": row["question"],
            "trajectory_text": row["trajectory_text"],
            "formatted_prompt": row["formatted_prompt"],
            "status": row["status"],
            "n_steps": int(row["n_steps"]),
            "n_tool_calls": int(row["n_tool_calls"]),
            "n_reasoning_steps": int(row["n_reasoning_steps"]),
            "run_id": row["run_id"],
        })
    _cache = rows
    return rows


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


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