app.py

"""Streamlit annotation dashboard.

Reads bundled data from ./data/, gates entry with password from env, lets each
annotator rate items pairwise (A / B / Tie + optional comment), pushes one
JSONL file per annotator to a private HF Dataset on every submit.
"""
from __future__ import annotations

import hashlib
import json
import os
import random
import time
from datetime import datetime, timezone
from pathlib import Path

import streamlit as st
from PIL import Image

import hf_io


APP_VERSION = "0.1.0"
DATA_DIR = Path(__file__).resolve().parent / "data"
BUNDLE_PATH = DATA_DIR / "annotation_data.json"
IMAGES_DIR = DATA_DIR / "images"
ANNOTATIONS_REPO = os.environ.get("ANNOTATIONS_REPO", hf_io.ANNOTATIONS_REPO_DEFAULT)


# --- boot ------------------------------------------------------------------

st.set_page_config(page_title="Bengali Headline Annotation", layout="wide")

# Bangla font (browser will fall back if Noto Sans Bengali isn't preinstalled).
st.markdown(
    """
    <link href="https://fonts.googleapis.com/css2?family=Noto+Sans+Bengali:wght@400;500;700&display=swap" rel="stylesheet">
    <style>
      .bn { font-family: 'Noto Sans Bengali', 'Hind Siliguri', sans-serif; line-height: 1.7; color: inherit; }
      .article-box { background: #fafafa; color: #111; border: 1px solid #ddd; padding: 1em;
                     border-radius: 6px; max-height: 380px; overflow-y: auto; }
      .headline-box { background: #ffffff; color: #111; border: 2px solid #ccc; padding: 1.2em;
                      border-radius: 8px; font-size: 1.4em; }
    </style>
    """,
    unsafe_allow_html=True,
)


@st.cache_data
def load_bundle() -> dict:
    if not BUNDLE_PATH.exists():
        st.error(f"Bundle not found at {BUNDLE_PATH}. Run build_annotation_data.py and redeploy.")
        st.stop()
    with BUNDLE_PATH.open() as f:
        return json.load(f)


def env_or_stop(key: str) -> str:
    val = os.environ.get(key)
    if not val:
        st.error(
            f"Missing required env var **{key}**. "
            f"Set it in Space → Settings → Repository secrets and reboot the Space."
        )
        st.stop()
    return val


def per_annotator_order(annotator: str, item_ids: list[str]) -> list[str]:
    h = hashlib.sha256(annotator.encode()).hexdigest()
    rng = random.Random(int(h, 16) % (2**64))
    shuffled = list(item_ids)
    rng.shuffle(shuffled)
    return shuffled


def _now_utc() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


# --- entry gate ------------------------------------------------------------

def render_entry_screen(expected_password: str) -> None:
    st.title("Bengali Headline Annotation")
    st.write("Please log in to begin.")
    st.caption("Please use only one browser tab while annotating — opening two tabs as the same user can overwrite each other's progress.")

    with st.form("entry"):
        name_raw = st.text_input("Your name (single word)", value="", autocomplete="off")
        password = st.text_input("Password", value="", type="password", autocomplete="off")
        submitted = st.form_submit_button("Start")

    if not submitted:
        return

    name = name_raw.strip().lower()
    if not name:
        st.error("Please enter your name.")
        return
    if any(c.isspace() for c in name):
        st.warning(
            f"Please use a **single word** for your name (no spaces). "
            f"We got: `{name_raw}`. Edit and try again."
        )
        return
    if password != expected_password:
        st.error("Incorrect password.")
        return

    st.session_state.annotator = name
    st.session_state.pending = []
    st.session_state.item_started_at = None
    st.rerun()


# --- annotation flow -------------------------------------------------------

def setup_session(bundle: dict, hf_token: str) -> None:
    """Compute remaining items for the current annotator. Cached in session_state."""
    if "remaining_ids" in st.session_state:
        return
    annotator = st.session_state.annotator
    all_items = bundle["items"]
    all_ids = [it["item_id"] for it in all_items]
    order = per_annotator_order(annotator, all_ids)

    try:
        existing = hf_io.read_annotator_file(ANNOTATIONS_REPO, annotator, token=hf_token)
    except Exception as e:
        st.error(f"Could not read existing annotations from HF: {e}")
        st.stop()

    done = {row["item_id"] for row in existing}
    remaining = [iid for iid in order if iid not in done]
    st.session_state.items_by_id = {it["item_id"]: it for it in all_items}
    st.session_state.all_order = order
    st.session_state.existing_rows = existing
    st.session_state.remaining_ids = remaining
    st.session_state.total = len(all_items)


def render_annotation_screen(hf_token: str) -> None:
    annotator = st.session_state.annotator
    remaining = st.session_state.remaining_ids
    total = st.session_state.total
    done_count = total - len(remaining)

    cols = st.columns([3, 1])
    with cols[0]:
        st.subheader(f"Annotator: `{annotator}`")
        st.progress(done_count / total if total else 0.0, text=f"{done_count} / {total} annotated")
    with cols[1]:
        if st.button("Log out"):
            for k in list(st.session_state.keys()):
                del st.session_state[k]
            st.rerun()

    if st.session_state.pending:
        st.warning(f"{len(st.session_state.pending)} ratings queued (last push failed). Will retry on next submit.")

    if not remaining:
        st.success(f"✓ Done — you rated {total} items. Thank you!")
        return

    item_id = remaining[0]
    item = st.session_state.items_by_id[item_id]

    if st.session_state.item_started_at is None or st.session_state.get("current_item_id") != item_id:
        st.session_state.item_started_at = time.time()
        st.session_state.current_item_id = item_id

    st.divider()
    st.caption(f"Source split: **{item['source_split']}** · item {done_count + 1} of {total}")

    left, right = st.columns([3, 2])
    with left:
        st.markdown("**Article**")
        st.markdown(f'<div class="bn article-box">{_html_escape(item["article"])}</div>', unsafe_allow_html=True)
    with right:
        if item.get("image_filename"):
            img_path = IMAGES_DIR / item["image_filename"]
            if img_path.exists():
                try:
                    st.image(Image.open(img_path), caption="Article image", use_container_width=True)
                except (OSError, Image.UnidentifiedImageError):
                    st.caption("[image could not be displayed]")
            else:
                st.caption("[image file missing]")
        else:
            st.caption("(no image for this split)")

    st.markdown("**Which headline is better?**")
    a_col, b_col = st.columns(2)
    with a_col:
        st.markdown("**A**")
        st.markdown(f'<div class="bn headline-box">{_html_escape(item["headline_left"])}</div>', unsafe_allow_html=True)
    with b_col:
        st.markdown("**B**")
        st.markdown(f'<div class="bn headline-box">{_html_escape(item["headline_right"])}</div>', unsafe_allow_html=True)

    st.write("")
    choice_label = st.radio(
        "Your choice",
        ["A is better", "B is better", "Tie"],
        index=None,
        horizontal=True,
        key=f"choice_{item_id}",
    )
    comment = st.text_area("Comment (optional)", key=f"comment_{item_id}", height=80)

    submit_disabled = choice_label is None
    if st.button("Submit", type="primary", disabled=submit_disabled):
        choice_map = {"A is better": "A_better", "B is better": "B_better", "Tie": "tie"}
        new_row = {
            "annotator": annotator,
            "item_id": item_id,
            "source_split": item["source_split"],
            "model_id": item["model_id"],
            "choice": choice_map[choice_label],
            "comment": comment.strip(),
            "timestamp_utc": _now_utc(),
            "app_version": APP_VERSION,
            "duration_ms": int((time.time() - st.session_state.item_started_at) * 1000),
        }
        all_rows = list(st.session_state.existing_rows) + list(st.session_state.pending) + [new_row]
        try:
            hf_io.upload_annotator_file(
                repo_id=ANNOTATIONS_REPO,
                annotator=annotator,
                rows=all_rows,
                token=hf_token,
                commit_message=f"{annotator}: item {done_count + 1}/{total}",
            )
        except Exception as e:
            st.session_state.pending.append(new_row)
            st.error(f"Push failed: {e}. Queued locally — will retry on next submit.")
            return
        st.session_state.existing_rows = all_rows
        st.session_state.pending = []
        st.session_state.remaining_ids = remaining[1:]
        st.session_state.item_started_at = None
        st.session_state.current_item_id = None
        st.rerun()


def _html_escape(s: str) -> str:
    return (
        s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
         .replace("\n", "<br>")
    )


# --- main ------------------------------------------------------------------

def main() -> None:
    bundle = load_bundle()
    hf_token = env_or_stop("HF_TOKEN")
    expected_password = env_or_stop("ANNOTATOR_PASSWORD")

    try:
        hf_io.ensure_dataset_exists(ANNOTATIONS_REPO, token=hf_token, private=True)
    except Exception as e:
        st.error(f"Could not access/create annotations dataset `{ANNOTATIONS_REPO}`: {e}")
        st.stop()

    if "annotator" not in st.session_state:
        render_entry_screen(expected_password)
        return

    setup_session(bundle, hf_token)
    render_annotation_screen(hf_token)


main()