data_studio/app.py

from __future__ import annotations

from typing import Any

import pandas as pd
import streamlit as st

from data_studio.auth import authenticate_user, count_users, get_active_user
from data_studio.config import ConfigError, StudioSettings
from data_studio.db import create_client, ensure_indexes, get_database
from data_studio.importer import import_source_samples
from data_studio.review import apply_review_action, claim_next_sample, queue_status_counts
from data_studio.sources import list_sources, seed_default_sources
from data_studio.utils import normalize_text, split_csv_tags, split_document_lines


@st.cache_resource
def load_settings() -> StudioSettings:
    return StudioSettings.from_env()


@st.cache_resource
def load_database() -> Any:
    settings = load_settings()
    client = create_client(settings)
    return get_database(client, settings)


def get_session_user() -> dict[str, Any] | None:
    username = st.session_state.get("auth_username")
    if not username:
        return None
    return {
        "username": username,
        "display_name": st.session_state.get("auth_display_name", username),
        "role": st.session_state.get("auth_role", "reviewer"),
    }


def set_session_user(user: dict[str, Any]) -> None:
    st.session_state["auth_username"] = user["username"]
    st.session_state["auth_display_name"] = user.get("display_name") or user["username"]
    st.session_state["auth_role"] = user.get("role", "reviewer")


def clear_session_user() -> None:
    for key in ("auth_username", "auth_display_name", "auth_role"):
        st.session_state.pop(key, None)


def render_login(db) -> None:
    st.title("Rachana Data Studio")
    st.caption("Invite-only access for approved reviewers.")
    if count_users(db) == 0:
        st.warning(
            "No users are configured yet. Create the first invite-only account with "
            "`python .\\scripts\\create_data_studio_user.py --username <name>`."
        )

    with st.form("login-form", clear_on_submit=False):
        username = st.text_input("Username")
        password = st.text_input("Password", type="password")
        submitted = st.form_submit_button("Log in", type="primary", use_container_width=True)

    if submitted:
        user = authenticate_user(db, username, password)
        if user is None:
            st.error("Invalid username or password, or the account is inactive.")
        else:
            set_session_user(user)
            st.rerun()


def require_authenticated_user(db) -> dict[str, Any] | None:
    session_user = get_session_user()
    if session_user is None:
        return None

    active_user = get_active_user(db, session_user["username"])
    if active_user is None:
        clear_session_user()
        st.warning("Your account is no longer active. Please log in again.")
        st.rerun()
    return {
        "username": active_user["username"],
        "display_name": active_user.get("display_name") or active_user["username"],
        "role": active_user.get("role", "reviewer"),
    }


def render_overview(db) -> None:
    counts = queue_status_counts(db)
    st.subheader("Queue overview")
    if counts:
        total_pending = sum(statuses.get("pending", 0) for statuses in counts.values())
        total_accepted = sum(statuses.get("accepted", 0) for statuses in counts.values())
        total_edited = sum(statuses.get("edited", 0) for statuses in counts.values())
        total_rejected = sum(statuses.get("rejected", 0) for statuses in counts.values())
        cards = st.columns(4)
        cards[0].metric("Pending", total_pending)
        cards[1].metric("Accepted", total_accepted)
        cards[2].metric("Edited", total_edited)
        cards[3].metric("Rejected", total_rejected)

        rows: list[dict[str, Any]] = []
        for queue_name, statuses in counts.items():
            row = {"queue_name": queue_name}
            row.update(statuses)
            rows.append(row)
        st.dataframe(pd.DataFrame(rows).fillna(0), width="stretch")
    else:
        st.info("No samples imported yet.")

    st.subheader("Source configs")
    source_rows = list_sources(db)
    if source_rows:
        st.dataframe(
            pd.DataFrame(
                [
                    {
                        "dataset_key": row["dataset_key"],
                        "queue_name": row["queue_name"],
                        "sample_type": row["sample_type"],
                        "hf_dataset": row["hf_dataset"],
                        "hf_config": row.get("hf_config"),
                        "hf_split": row["hf_split"],
                    }
                    for row in source_rows
                ]
            ),
            width="stretch",
        )
    else:
        st.warning("No source configs seeded yet.")


def render_sources(db) -> None:
    st.subheader("Source setup")
    st.caption("Seed the default source configurations for the active review queues.")
    if st.button("Seed default sources", type="primary"):
        inserted = seed_default_sources(db)
        st.success(f"Source config seed completed. New records inserted: {inserted}.")
        st.rerun()

    source_rows = list_sources(db)
    if source_rows:
        for row in source_rows:
            with st.container(border=True):
                top = st.columns([2, 2, 2, 1])
                top[0].markdown(f"**{row['name']}**")
                top[1].caption(f"Queue: `{row['queue_name']}`")
                top[2].caption(f"Type: `{row['sample_type']}`")
                top[3].caption("Enabled" if row.get("enabled", True) else "Disabled")
                st.caption(f"Dataset key: `{row['dataset_key']}`")
                st.caption(f"HF dataset: `{row['hf_dataset']}`")
                if row.get("description"):
                    st.write(row["description"])


def render_import(db, settings: StudioSettings) -> None:
    st.subheader("Import Hugging Face source data")
    sources = list_sources(db)
    if not sources:
        st.warning("Seed source configs first.")
        return

    dataset_keys = [source["dataset_key"] for source in sources if source.get("enabled", True)]
    selected = st.selectbox("Source", dataset_keys)
    selected_source = next(source for source in sources if source["dataset_key"] == selected)
    current_counts = queue_status_counts(db).get(selected_source["queue_name"], {})

    info_cols = st.columns(4)
    info_cols[0].metric("Queue", selected_source["queue_name"])
    info_cols[1].metric("Pending now", current_counts.get("pending", 0))
    info_cols[2].metric("Accepted", current_counts.get("accepted", 0))
    info_cols[3].metric("Edited", current_counts.get("edited", 0))
    st.caption(f"HF dataset: `{selected_source['hf_dataset']}`")

    max_samples = st.number_input("Max samples to import", min_value=10, max_value=50000, value=1000, step=10)

    quick_cols = st.columns(4)
    if quick_cols[0].button("Import 25", use_container_width=True):
        max_samples = 25
        st.session_state["import_now"] = True
    if quick_cols[1].button("Import 100", use_container_width=True):
        max_samples = 100
        st.session_state["import_now"] = True
    if quick_cols[2].button("Import 500", use_container_width=True):
        max_samples = 500
        st.session_state["import_now"] = True
    manual_import = quick_cols[3].button("Import custom", type="primary", use_container_width=True)

    if manual_import or st.session_state.pop("import_now", False):
        with st.spinner("Importing samples from Hugging Face..."):
            stats = import_source_samples(db, settings, selected, int(max_samples))
        st.success("Import complete.")
        st.json(stats)


def render_document_editor(sample: dict[str, Any]) -> dict[str, Any]:
    sample_key = sample["sample_id"]
    with st.expander("Original document", expanded=False):
        st.text_area("Original", value=sample["payload"]["text"], height=220, disabled=True)

    current_lines = split_document_lines(sample["review"]["current_text"])
    line_df = pd.DataFrame(
        {
            "keep": [True] * len(current_lines),
            "line_no": list(range(1, len(current_lines) + 1)),
            "text": current_lines,
        }
    )
    edited_df = st.data_editor(
        line_df,
        width="stretch",
        hide_index=True,
        num_rows="fixed",
        column_config={
            "keep": st.column_config.CheckboxColumn("Keep"),
            "line_no": st.column_config.NumberColumn("Line", disabled=True),
            "text": st.column_config.TextColumn("Text", width="large"),
        },
        disabled=["line_no"],
        key=f"{sample['sample_id']}-line-editor",
    )

    kept_lines = [
        normalize_text(str(row["text"]))
        for _, row in edited_df.iterrows()
        if bool(row["keep"]) and normalize_text(str(row["text"]))
    ]

    added_lines_raw = st.text_area(
        "Add new lines (one per line)",
        value="",
        height=100,
        key=f"{sample_key}-added-lines",
    )
    added_lines = [normalize_text(line) for line in added_lines_raw.splitlines()]
    kept_lines.extend(line for line in added_lines if line)

    current_text = "\n".join(kept_lines)
    st.caption(f"Kept lines: {len(kept_lines)}")
    st.text_area("Final cleaned text preview", value=current_text, height=180, disabled=True)
    return {"current_text": current_text}


def render_translation_editor(sample: dict[str, Any]) -> dict[str, Any]:
    col1, col2 = st.columns(2)
    with col1:
        st.markdown("**Original source**")
        st.text_area("Original source text", value=sample["payload"]["source_text"], height=180, disabled=True)
        source_text = st.text_area(
            "Reviewed source text",
            value=sample["review"]["current_pair"]["source_text"],
            height=180,
        )
    with col2:
        st.markdown("**Original target**")
        st.text_area("Original target text", value=sample["payload"]["target_text"], height=180, disabled=True)
        target_text = st.text_area(
            "Reviewed target text",
            value=sample["review"]["current_pair"]["target_text"],
            height=180,
        )
    return {"current_pair": {"source_text": source_text, "target_text": target_text}}


def render_transliteration_editor(sample: dict[str, Any]) -> dict[str, Any]:
    col1, col2 = st.columns(2)
    with col1:
        st.text_input("Original native text", value=sample["payload"]["native_text"], disabled=True)
        native_text = st.text_input(
            "Reviewed native text",
            value=sample["review"]["current_transliteration"]["native_text"],
        )
    with col2:
        st.text_input("Original Latin text", value=sample["payload"]["latin_text"], disabled=True)
        latin_text = st.text_input(
            "Reviewed Latin text",
            value=sample["review"]["current_transliteration"]["latin_text"],
        )
    return {"current_transliteration": {"native_text": native_text, "latin_text": latin_text}}


def render_review(db, user: dict[str, Any]) -> None:
    st.subheader("Review queue")
    counts = queue_status_counts(db)
    sources = [source for source in list_sources(db) if source.get("enabled", True)]
    dataset_options = ["All active datasets"] + [source["dataset_key"] for source in sources]
    selected_dataset = st.selectbox("Dataset", dataset_options)
    queue_name = "all"
    queue_stats = counts.get("pure_telugu", {})
    stats_cols = st.columns(4)
    stats_cols[0].metric("Pending", queue_stats.get("pending", 0))
    stats_cols[1].metric("Accepted", queue_stats.get("accepted", 0))
    stats_cols[2].metric("Edited", queue_stats.get("edited", 0))
    stats_cols[3].metric("Rejected", queue_stats.get("rejected", 0))

    sample = claim_next_sample(
        db,
        queue_name,
        user["username"],
        None if selected_dataset == "All active datasets" else selected_dataset,
    )
    if sample is None:
        st.info("No pending samples for the selected pool.")
        return

    st.caption(f"Sample ID: {sample['sample_id']}")
    meta_cols = st.columns(4)
    meta_cols[0].markdown(f"**Dataset**\n\n`{sample['source']['dataset_key']}`")
    meta_cols[1].markdown(f"**Record ID**\n\n`{sample['source'].get('source_record_id')}`")
    meta_cols[2].markdown(f"**Status**\n\n`{sample['status']}`")
    meta_cols[3].markdown(f"**Type**\n\n`{sample['sample_type']}`")
    if sample["source"].get("source_title"):
        st.caption(f"Title: {sample['source']['source_title']}")

    with st.form(f"review-{sample['sample_id']}"):
        if sample["sample_type"] == "document":
            cleaned_payload = render_document_editor(sample)
        elif sample["sample_type"] == "translation_pair":
            cleaned_payload = render_translation_editor(sample)
        else:
            cleaned_payload = render_transliteration_editor(sample)

        with st.expander("Tags and review notes", expanded=True):
            keyword_text = st.text_input(
                "Keywords / tags (comma separated)",
                value=", ".join(sample["review"]["quality_tags"]),
            )
            task_tag_text = st.text_input(
                "Task tags (comma separated)",
                value=", ".join(sample["review"]["task_tags"]),
            )
            notes = st.text_area("Reviewer notes", value=sample["review"].get("notes", ""), height=100)
            reason = st.text_input("Action reason")
            confidence = st.slider("Reviewer confidence", min_value=1, max_value=5, value=4)

        col1, col2, col3, col4 = st.columns(4)
        accept = col1.form_submit_button("Accept", type="primary", use_container_width=True)
        edit = col2.form_submit_button("Save Edit", use_container_width=True)
        reject = col3.form_submit_button("Reject", use_container_width=True)
        skip = col4.form_submit_button("Skip", use_container_width=True)

    action = None
    if accept:
        action = "accept"
    elif edit:
        action = "edit"
    elif reject:
        action = "reject"
    elif skip:
        action = "skip"

    if action is not None:
        apply_review_action(
            db=db,
            sample_id=sample["sample_id"],
            action=action,
            reviewer=user["username"],
            expected_version=sample["version"],
            cleaned_payload=cleaned_payload,
            tags=split_csv_tags(keyword_text),
            quality_tags=split_csv_tags(keyword_text),
            task_tags=split_csv_tags(task_tag_text),
            notes=notes,
            confidence=confidence,
            reason=reason,
        )
        st.success(f"Action recorded: {action}")
        st.rerun()


def main() -> None:
    st.set_page_config(page_title="Rachana Data Studio", layout="wide")

    try:
        settings = load_settings()
        db = load_database()
    except ConfigError as exc:
        st.error(str(exc))
        st.stop()

    ensure_indexes(db)
    user = require_authenticated_user(db)
    if user is None:
        render_login(db)
        return

    st.title("Rachana Data Studio")
    st.caption("HF-native corpus artifacts with MongoDB-backed review state.")

    st.sidebar.header("Workspace")
    section = st.sidebar.radio("Section", ["Overview", "Sources", "Import", "Review"])
    st.sidebar.caption(f"MongoDB DB: `{settings.mongodb_db_name}`")
    st.sidebar.caption(f"Signed in as: `{user['display_name']}`")
    st.sidebar.caption(f"Role: `{user['role']}`")
    if st.sidebar.button("Log out", use_container_width=True):
        clear_session_user()
        st.rerun()

    if section == "Overview":
        render_overview(db)
    elif section == "Sources":
        render_sources(db)
    elif section == "Import":
        render_import(db, settings)
    else:
        render_review(db, user)


if __name__ == "__main__":
    main()