sheets.py

"""Persist eval scenarios in a Google Sheet (the source of truth on HF Spaces,
whose filesystem is ephemeral).

One row per scenario; the ``expect`` block is flattened into typed columns. The
app passes the same service-account creds dict it already uses for Google Docs
(needs the spreadsheets scope) plus SCENARIOS_SHEET_ID.
"""
from __future__ import annotations

from typing import Any

# Column order written to row 1. Loading maps by header name, so reordering
# columns in the sheet is safe; renaming them is not.
COLUMNS = [
    "id", "turns", "use_tools",
    "escalate", "format_ok", "no_prompt_leak", "tools_any", "tools_none", "judge",
]

SCOPES = ["https://www.googleapis.com/auth/spreadsheets"]


def _service(creds_dict: dict):
    from google.oauth2.service_account import Credentials
    from googleapiclient.discovery import build

    creds = Credentials.from_service_account_info(creds_dict, scopes=SCOPES)
    return build("sheets", "v4", credentials=creds, cache_discovery=False)


def _first_tab(svc, sheet_id: str) -> str:
    meta = svc.spreadsheets().get(spreadsheetId=sheet_id).execute()
    sheets = meta.get("sheets", [])
    if not sheets:
        return "Sheet1"
    return sheets[0]["properties"]["title"]


def _to_bool(s: str):
    v = (s or "").strip().lower()
    if v in ("true", "1", "yes", "да", "y"):
        return True
    if v in ("false", "0", "no", "нет", "n"):
        return False
    return None


def _b(v) -> str:
    if v is True:
        return "TRUE"
    if v is False:
        return "FALSE"
    return ""


def _row_to_scenario(row: list[str], header: list[str]) -> dict:
    d = {h: (row[i] if i < len(row) else "") for i, h in enumerate(header)}
    sc: dict[str, Any] = {
        "id": (d.get("id", "") or "").strip(),
    }
    sc["turns"] = [t for t in (d.get("turns", "") or "").splitlines() if t.strip()]

    ut = _to_bool(d.get("use_tools", ""))
    if ut is not None:
        sc["use_tools"] = ut

    expect: dict[str, Any] = {}
    for key in ("escalate", "format_ok"):
        b = _to_bool(d.get(key, ""))
        if b is not None:
            expect[key] = b
    for key in ("no_prompt_leak", "tools_none"):
        if _to_bool(d.get(key, "")) is True:
            expect[key] = True
    tools_any = [t.strip() for t in (d.get("tools_any", "") or "").split(",") if t.strip()]
    if tools_any:
        expect["tools_any"] = tools_any
    judge = (d.get("judge", "") or "").strip()
    if judge:
        expect["judge"] = judge
    sc["expect"] = expect
    return sc


def _scenario_to_row(sc: dict) -> list[str]:
    expect = sc.get("expect", {}) or {}
    turns = sc.get("turns") or ([sc["message"]] if sc.get("message") else [])
    return [
        sc.get("id", "") or "",
        "\n".join(turns),
        _b(sc.get("use_tools")) if "use_tools" in sc else "",
        _b(expect.get("escalate")) if "escalate" in expect else "",
        _b(expect.get("format_ok")) if "format_ok" in expect else "",
        "TRUE" if expect.get("no_prompt_leak") else "",
        ", ".join(expect.get("tools_any", []) or []),
        "TRUE" if expect.get("tools_none") else "",
        expect.get("judge", "") or "",
    ]


def load_scenarios_from_sheet(creds_dict: dict, sheet_id: str) -> list[dict]:
    svc = _service(creds_dict)
    tab = _first_tab(svc, sheet_id)
    resp = (
        svc.spreadsheets()
        .values()
        .get(spreadsheetId=sheet_id, range=tab)
        .execute()
    )
    values = resp.get("values", [])
    if not values:
        return []
    header = [h.strip() for h in values[0]]
    out = []
    for row in values[1:]:
        sc = _row_to_scenario(row, header)
        if sc["id"]:
            out.append(sc)
    return out


def save_scenarios_to_sheet(creds_dict: dict, sheet_id: str, scenarios: list[dict]) -> None:
    svc = _service(creds_dict)
    tab = _first_tab(svc, sheet_id)
    rows = [COLUMNS] + [_scenario_to_row(sc) for sc in scenarios]
    svc.spreadsheets().values().clear(spreadsheetId=sheet_id, range=tab).execute()
    (
        svc.spreadsheets()
        .values()
        .update(
            spreadsheetId=sheet_id,
            range=f"{tab}!A1",
            valueInputOption="RAW",
            body={"values": rows},
        )
        .execute()
    )