github_storage.py

from __future__ import annotations

import base64
import csv
import io
import json
import logging
from typing import Any

import requests
import streamlit as st

logger = logging.getLogger(__name__)

FIELDNAMES = [
    "pair_index",
    "reviewer_id",
    "clinically_plausible",
    "pathology_preserved",
    "quality_score",
    "comments",
    "timestamp",
]


def _get_config() -> dict[str, str]:
    import os
    if os.environ.get("GITHUB_TOKEN"):
        return {
            "token": os.environ["GITHUB_TOKEN"],
            "repo": os.environ.get("GITHUB_REPO", ""),
            "branch": os.environ.get("GITHUB_BRANCH", "validation-data"),
            "data_dir": os.environ.get("GITHUB_DATA_DIR", "validation_data"),
        }
    return {
        "token": st.secrets["github"]["token"],
        "repo": st.secrets["github"]["repo"],
        "branch": st.secrets["github"].get("branch", "validation-data"),
        "data_dir": st.secrets["github"].get("data_dir", "validation_data"),
    }


def _headers(token: str) -> dict[str, str]:
    return {
        "Authorization": f"Bearer {token}",
        "Accept": "application/vnd.github.v3+json",
        "X-GitHub-Api-Version": "2022-11-28",
    }


def _get_file(repo: str, path: str, token: str, branch: str) -> tuple[str | None, str | None]:
    url = f"https://api.github.com/repos/{repo}/contents/{path}"
    resp = requests.get(url, headers=_headers(token), params={"ref": branch}, timeout=15)
    if resp.status_code == 404:
        return None, None
    resp.raise_for_status()
    data = resp.json()
    content = base64.b64decode(data["content"]).decode("utf-8")
    return content, data["sha"]


def _put_file(
    repo: str,
    path: str,
    content: str,
    token: str,
    branch: str,
    sha: str | None,
    message: str,
) -> None:
    url = f"https://api.github.com/repos/{repo}/contents/{path}"
    body: dict[str, Any] = {
        "message": message,
        "content": base64.b64encode(content.encode("utf-8")).decode("ascii"),
        "branch": branch,
    }
    if sha:
        body["sha"] = sha
    resp = requests.put(url, headers=_headers(token), json=body, timeout=15)
    resp.raise_for_status()


def _ensure_branch(repo: str, token: str, branch: str) -> None:
    url = f"https://api.github.com/repos/{repo}/branches/{branch}"
    resp = requests.get(url, headers=_headers(token), timeout=15)
    if resp.status_code == 200:
        return

    main_url = f"https://api.github.com/repos/{repo}/git/refs/heads/main"
    main_resp = requests.get(main_url, headers=_headers(token), timeout=15)
    main_resp.raise_for_status()
    main_sha = main_resp.json()["object"]["sha"]

    create_url = f"https://api.github.com/repos/{repo}/git/refs"
    create_body = {"ref": f"refs/heads/{branch}", "sha": main_sha}
    create_resp = requests.post(create_url, headers=_headers(token), json=create_body, timeout=15)
    if create_resp.status_code not in (200, 201, 422):
        create_resp.raise_for_status()


def is_configured() -> bool:
    import os
    if os.environ.get("GITHUB_TOKEN"):
        return True
    try:
        _get_config()
        return True
    except (KeyError, FileNotFoundError):
        return False


def load_pairs() -> list[dict]:
    cfg = _get_config()
    content, _ = _get_file(cfg["repo"], f"{cfg['data_dir']}/pairs.json", cfg["token"], cfg["branch"])
    if content is None:
        return []
    return json.loads(content)


def load_existing_reviews(reviewer_id: str) -> dict[int, dict]:
    cfg = _get_config()
    path = f"{cfg['data_dir']}/reviews_{reviewer_id}.csv"
    content, _ = _get_file(cfg["repo"], path, cfg["token"], cfg["branch"])
    if content is None:
        return {}
    reviews: dict[int, dict] = {}
    reader = csv.DictReader(io.StringIO(content))
    for row in reader:
        reviews[int(row["pair_index"])] = row
    return reviews


def save_review(reviewer_id: str, pair_index: int, review: dict) -> None:
    cfg = _get_config()
    _ensure_branch(cfg["repo"], cfg["token"], cfg["branch"])

    path = f"{cfg['data_dir']}/reviews_{reviewer_id}.csv"
    content, sha = _get_file(cfg["repo"], path, cfg["token"], cfg["branch"])

    existing: dict[int, dict] = {}
    if content:
        reader = csv.DictReader(io.StringIO(content))
        for row in reader:
            existing[int(row["pair_index"])] = row

    existing[pair_index] = review

    output = io.StringIO()
    writer = csv.DictWriter(output, fieldnames=FIELDNAMES)
    writer.writeheader()
    for idx in sorted(existing.keys()):
        writer.writerow(existing[idx])

    _put_file(
        cfg["repo"],
        path,
        output.getvalue(),
        cfg["token"],
        cfg["branch"],
        sha,
        f"Review by {reviewer_id}: pair {pair_index}",
    )


def append_audit_log(reviewer_id: str, pair_index: int, action: str) -> None:
    import time

    cfg = _get_config()
    _ensure_branch(cfg["repo"], cfg["token"], cfg["branch"])

    path = f"{cfg['data_dir']}/audit_log.csv"
    content, sha = _get_file(cfg["repo"], path, cfg["token"], cfg["branch"])

    if content is None:
        content = "timestamp,reviewer_id,pair_index,action\n"
        sha = None

    timestamp = time.strftime("%Y-%m-%d %H:%M:%S UTC", time.gmtime())
    content += f"{timestamp},{reviewer_id},{pair_index},{action}\n"

    _put_file(
        cfg["repo"],
        path,
        content,
        cfg["token"],
        cfg["branch"],
        sha,
        f"Audit: {action} by {reviewer_id}",
    )