"""Romanian Cultural VLM Benchmark — annotation app (v1).

Run:
    streamlit run app.py
"""
from __future__ import annotations

import os
import urllib.parse

import streamlit as st

import storage

DEFAULT_TYPE = "mcq"
MAX_SLOTS = 8
JOIN_CODE = os.environ.get("JOIN_CODE", "").strip()
ADMIN_USERNAMES = {
    u.strip() for u in os.environ.get("ADMIN_USERNAMES", "").split(",") if u.strip()
}
ANS_NONE = "—"  # string sentinel for "no answer selected" in the radio widget


# ---------- session helpers ----------

def _init_state() -> None:
    st.session_state.setdefault("username", None)
    st.session_state.setdefault("current_image", None)
    st.session_state.setdefault("qa_slots", None)
    st.session_state.setdefault("skipped_ids", set())


def _blank_slot(slot_idx: int, qa_type: str = DEFAULT_TYPE) -> dict:
    return {
        "id": storage.new_qa_id(),
        "slot": slot_idx,
        "qa_type": qa_type,
        "is_cultural": False,
        "question_ro": "",
        "answer_ro": "",
        "options_ro": ["", "", "", ""],
        "answer_idx": None,
        "supporting_paragraph": "",
    }


def _reindex_slots(slots: list[dict]) -> None:
    for i, s in enumerate(slots):
        s["slot"] = i


def _coerce_slot(raw: dict, slot_idx: int) -> dict:
    """Defensive: build a clean slot dict from a (possibly stale or wrongly
    typed) draft record. Replaces any None / missing / wrong-type fields
    with sane defaults so the form never sees a None it can't handle."""
    s = _blank_slot(slot_idx)
    if isinstance(raw.get("id"), str):
        s["id"] = raw["id"]
    if raw.get("qa_type") in ("mcq", "open"):
        s["qa_type"] = raw["qa_type"]
    for key in ("question_ro", "answer_ro", "supporting_paragraph"):
        if isinstance(raw.get(key), str):
            s[key] = raw[key]
    if isinstance(raw.get("is_cultural"), bool):
        s["is_cultural"] = raw["is_cultural"]
    opts = raw.get("options_ro")
    if isinstance(opts, list) and len(opts) == 4:
        s["options_ro"] = [str(x) if x is not None else "" for x in opts]
    ai = raw.get("answer_idx")
    if isinstance(ai, int) and 0 <= ai < 4:
        s["answer_idx"] = ai
    return s


def _slot_to_draft(slot: dict, image: dict, username: str) -> dict:
    """Augment a raw slot with image/user metadata for storage."""
    return {
        **slot,
        "image_id": image["filename"],
        "image_url": image["url"],
        "author_username": username,
    }


def _load_image_for_user(username: str) -> None:
    """Pick the next image; hydrate slots from drafts if any, else a single blank slot."""
    image = storage.next_image(username, exclude=st.session_state.skipped_ids)
    if image is None:
        st.session_state.current_image = None
        st.session_state.qa_slots = None
        return

    storage.acquire_lock(username, image["filename"])
    drafts = storage.get_drafts(username, image["filename"])
    if drafts:
        slots = [_coerce_slot(d, i) for i, d in enumerate(drafts)]
    else:
        slots = [_blank_slot(0)]

    st.session_state.current_image = image
    st.session_state.qa_slots = slots


def _on_add_question(username: str) -> None:
    """on_click callback — runs after Streamlit has committed all
    widget values to session_state, so typed text is captured."""
    slots = st.session_state.get("qa_slots") or []
    if len(slots) >= MAX_SLOTS:
        return
    slots = list(slots)  # new list to make sure Streamlit sees the change
    slots.append(_blank_slot(len(slots)))
    st.session_state.qa_slots = slots
    _autosave(username)


def _on_submit_image(username: str) -> None:
    image = st.session_state.current_image
    slots = st.session_state.get("qa_slots") or []
    if image is None or not slots:
        return
    records = [_build_record_from_session(s, image, username) for s in slots]
    storage.submit_image(records)
    _load_image_for_user(username)


def _on_skip_image(username: str) -> None:
    image = st.session_state.current_image
    if image is None:
        return
    _autosave(username)
    st.session_state.skipped_ids.add(image["filename"])
    storage.release_lock(image["filename"])
    _load_image_for_user(username)


def _on_save_draft(username: str) -> None:
    _autosave(username)


def _on_logout() -> None:
    image = st.session_state.current_image
    if image:
        _autosave(st.session_state.username)
        storage.release_lock(image["filename"])
    for k in ("username", "current_image", "qa_slots"):
        st.session_state[k] = None


def _build_record_from_session(slot: dict, image: dict, username: str) -> dict:
    """Like _build_record but reads field values from session_state to
    avoid stale slot-dict reads during button on_click callbacks."""
    sid = slot["id"]
    qa_type = st.session_state.get(f"type_{sid}", slot["qa_type"])
    question_ro = st.session_state.get(f"q_{sid}", slot["question_ro"]).strip()
    supporting = (
        st.session_state.get(f"sup_{sid}", slot["supporting_paragraph"]).strip()
        or None
    )
    is_cultural = bool(
        st.session_state.get(f"cult_{sid}", slot.get("is_cultural", False))
    )
    base = {
        "id": slot["id"],
        "image_id": image["filename"],
        "image_source": "wikimedia_commons",
        "image_url": image["url"],
        "qa_type": qa_type,
        "is_cultural": is_cultural,
        "question_ro": question_ro,
        "supporting_paragraph": supporting,
        "author_username": username,
        "concept_tags": [],
        "sensitivity_flags": [],
    }
    if qa_type == "mcq":
        opts = [
            st.session_state.get(f"opt_{sid}_{i}", slot["options_ro"][i]).strip()
            for i in range(4)
        ]
        base["options_ro"] = opts
        ans = st.session_state.get(f"ans_{sid}")
        idx = None if ans in (None, ANS_NONE) else ans
        base["answer_idx"] = idx
        base["answer_ro"] = opts[idx] if idx is not None and opts else ""
    else:
        base["answer_ro"] = st.session_state.get(
            f"a_{sid}", slot["answer_ro"]
        ).strip()
    return base


def _autosave(username: str) -> None:
    """Save current slot state as drafts. Reads directly from
    st.session_state so we always capture the latest typed values,
    not stale slot-dict mutations from the previous script run."""
    image = st.session_state.current_image
    slots = st.session_state.get("qa_slots")
    if image is None or not slots:
        return
    records = []
    for slot in slots:
        sid = slot["id"]
        ans = st.session_state.get(f"ans_{sid}")
        rec = {
            **slot,
            "qa_type": st.session_state.get(f"type_{sid}", slot["qa_type"]),
            "is_cultural": bool(
                st.session_state.get(f"cult_{sid}", slot.get("is_cultural", False))
            ),
            "question_ro": st.session_state.get(f"q_{sid}", slot["question_ro"]),
            "answer_ro": st.session_state.get(f"a_{sid}", slot["answer_ro"]),
            "supporting_paragraph": st.session_state.get(
                f"sup_{sid}", slot["supporting_paragraph"]
            ),
            "answer_idx": None if ans in (None, ANS_NONE) else ans,
            "options_ro": [
                st.session_state.get(f"opt_{sid}_{i}", slot["options_ro"][i])
                for i in range(4)
            ],
            "image_id": image["filename"],
            "image_url": image["url"],
            "author_username": username,
        }
        records.append(rec)
    storage.save_drafts(records)


# ---------- validation ----------

def _validate(slot: dict) -> str | None:
    if not slot["question_ro"].strip():
        return "Întrebarea este goală."
    if slot["qa_type"] == "mcq":
        opts = [o.strip() for o in slot["options_ro"]]
        if any(not o for o in opts):
            return "Toate cele 4 opțiuni MCQ trebuie completate."
        if len(set(opts)) != 4:
            return "Opțiunile MCQ trebuie să fie distincte."
        if slot.get("answer_idx") is None:
            return "Selectează răspunsul corect."
    else:
        if not slot["answer_ro"].strip():
            return "Răspunsul de referință este gol."
    return None


def _validate_all(slots: list[dict]) -> list[str]:
    errors: list[str] = []
    for s in slots:
        err = _validate(s)
        if err:
            errors.append(f"Q{s['slot'] + 1}: {err}")
    return errors


# ---------- record building ----------

def _build_record(slot: dict, image: dict, username: str) -> dict:
    base = {
        "id": slot["id"],
        "image_id": image["filename"],
        "image_source": "wikimedia_commons",
        "image_url": image["url"],
        "qa_type": slot["qa_type"],
        "is_cultural": bool(slot.get("is_cultural", False)),
        "question_ro": slot["question_ro"].strip(),
        "supporting_paragraph": slot["supporting_paragraph"].strip() or None,
        "author_username": username,
        "concept_tags": [],
        "sensitivity_flags": [],
    }
    if slot["qa_type"] == "mcq":
        opts = [o.strip() for o in slot["options_ro"]]
        base["options_ro"] = opts
        base["answer_idx"] = slot["answer_idx"]
        base["answer_ro"] = opts[slot["answer_idx"]] if opts else ""
    else:
        base["answer_ro"] = slot["answer_ro"].strip()
    return base


# ---------- UI ----------

STICKY_CSS = """
<style>
/* Don't let flex items stretch — needed for sticky to engage on the column. */
[data-testid="stHorizontalBlock"] {
    align-items: flex-start !important;
}

/* Sticky image column: any column containing an <img>. Nested MCQ
   option columns contain only text inputs, so they're unaffected. */
[data-testid="stColumn"]:has(img),
[data-testid="column"]:has(img),
.stColumn:has(img) {
    position: sticky !important;
    top: 1rem !important;
    align-self: flex-start !important;
    z-index: 10;
}
</style>
"""


def _login_view() -> None:
    st.title("Romanian Cultural VLM Benchmark — Adnotare")
    st.write("Introdu un nume de utilizator pentru a începe.")
    with st.form("login"):
        name = st.text_input("Nume utilizator").strip()
        code = ""
        if JOIN_CODE:
            code = st.text_input("Cod de acces", type="password").strip()
        submitted = st.form_submit_button("Intră")
        if submitted:
            if not name:
                st.error("Nume utilizator obligatoriu.")
                return
            if JOIN_CODE and code != JOIN_CODE:
                st.error("Cod de acces incorect.")
                return
            st.session_state.username = name
            st.session_state.skipped_ids = set()
            _load_image_for_user(name)
            st.rerun()


def _commons_page_url(filename: str) -> str:
    """Wikimedia Commons file page (description + full license + source)."""
    slug = urllib.parse.quote(filename.replace(" ", "_"))
    return f"https://commons.wikimedia.org/wiki/File:{slug}"


def _image_panel(image: dict) -> None:
    st.image(image["url"], use_container_width=True)
    st.caption(image["filename"])
    with st.expander("Metadata imagine", expanded=False):
        st.markdown(f"**Label RO:** {image.get('label_ro') or '—'}")
        st.markdown(f"**Label EN:** {image.get('label_en') or '—'}")
        wiki_cats = image.get("wikimedia_categories") or image.get("categories") or []
        st.markdown(f"**Categorii Wikimedia:** {', '.join(wiki_cats) or '—'}")
        concept_cats = image.get("concept_categories") or []
        if concept_cats:
            st.markdown(f"**Categorii concept:** {', '.join(concept_cats)}")
        st.markdown(f"**Licență:** {image.get('license') or '—'}")
        st.markdown(f"**Atribuire:** {image.get('attribution') or '—'}", unsafe_allow_html=True)
        st.markdown(f"**Score cultural:** {image.get('score', '—')}")
        st.markdown(
            f"**Pagină Commons:** [{image['filename']}]({_commons_page_url(image['filename'])})"
        )
        st.markdown(f"**Imagine originală:** [link direct]({image['url']})")
        if image.get("qid"):
            st.markdown(f"**QID:** [{image['qid']}](https://www.wikidata.org/wiki/{image['qid']})")


def _init_widget(key: str, default) -> None:
    """Seed session_state for a widget key before its first render so we
    can use the canonical key-only widget pattern (no `value=` argument
    passed to the widget — that combo is fragile across reruns)."""
    if key not in st.session_state:
        st.session_state[key] = default


def _slot_form(slot: dict, can_remove: bool) -> bool:
    """Render one Q/A slot. Returns True if user requested removal."""
    sid = slot["id"]
    type_key = f"type_{sid}"
    cult_key = f"cult_{sid}"
    q_key = f"q_{sid}"
    a_key = f"a_{sid}"
    sup_key = f"sup_{sid}"
    ans_key = f"ans_{sid}"
    opt_keys = [f"opt_{sid}_{i}" for i in range(4)]

    _init_widget(type_key, slot["qa_type"])
    _init_widget(cult_key, slot["is_cultural"])
    _init_widget(q_key, slot["question_ro"])
    _init_widget(a_key, slot["answer_ro"])
    _init_widget(sup_key, slot["supporting_paragraph"])
    # ans uses a string sentinel inside session_state to avoid None-in-options pitfalls
    _init_widget(
        ans_key,
        ANS_NONE if slot["answer_idx"] is None else slot["answer_idx"],
    )
    for i, k in enumerate(opt_keys):
        _init_widget(k, slot["options_ro"][i])

    n = slot["slot"] + 1
    header_cols = st.columns([5, 1])
    with header_cols[0]:
        st.markdown(f"### Q{n}")
    with header_cols[1]:
        if can_remove and st.button("Șterge", key=f"rm_{sid}"):
            return True

    slot["qa_type"] = st.radio(
        "Tip",
        options=["mcq", "open"],
        format_func=lambda x: "MCQ (alegere multiplă)" if x == "mcq" else "Deschisă",
        key=type_key,
        horizontal=True,
    )
    slot["is_cultural"] = st.checkbox("Întrebare culturală", key=cult_key)
    slot["question_ro"] = st.text_area(
        "Întrebare (RO)",
        key=q_key,
        height=80,
    )

    if slot["qa_type"] == "mcq":
        cols = st.columns(2)
        new_opts = list(slot["options_ro"])
        for i in range(4):
            with cols[i % 2]:
                new_opts[i] = st.text_input(
                    f"Opțiunea {i + 1}",
                    key=opt_keys[i],
                )
        slot["options_ro"] = new_opts
        raw_ans = st.radio(
            "Răspuns corect",
            options=[ANS_NONE, 0, 1, 2, 3],
            format_func=lambda i: "(neselectat)" if i == ANS_NONE else f"Opțiunea {i + 1}",
            key=ans_key,
            horizontal=True,
        )
        slot["answer_idx"] = None if raw_ans == ANS_NONE else raw_ans
    else:
        slot["answer_ro"] = st.text_area(
            "Răspuns de referință (RO)",
            key=a_key,
            height=80,
        )

    slot["supporting_paragraph"] = st.text_area(
        "Context cultural (opțional)",
        key=sup_key,
        height=68,
    )
    st.divider()
    return False


def _annotation_view() -> None:
    username = st.session_state.username
    image = st.session_state.current_image

    with st.sidebar:
        st.markdown(f"**Utilizator:** `{username}`")
        stats = storage.user_stats(username)
        st.metric("Imagini trimise", stats["images"])
        st.metric("Q/A trimise", stats["qa_pairs"])
        st.divider()
        st.button("Ieșire", on_click=_on_logout)

    if image is None:
        st.success("Coadă goală — toate imaginile curate au fost adnotate.")
        return

    left, right = st.columns([1, 2], gap="large")
    with left:
        _image_panel(image)
        c1, c2 = st.columns(2)
        c1.button(
            "Sari peste",
            use_container_width=True,
            on_click=_on_skip_image,
            args=(username,),
        )
        c2.button(
            "Salvează draft",
            type="secondary",
            use_container_width=True,
            on_click=_on_save_draft,
            args=(username,),
        )

    with right:
        slots = st.session_state.qa_slots
        remove_idx: int | None = None
        for i, slot in enumerate(slots):
            if _slot_form(slot, can_remove=len(slots) > 1):
                remove_idx = i
        if remove_idx is not None:
            slots.pop(remove_idx)
            _reindex_slots(slots)
            _autosave(username)
            st.rerun()

        add_disabled = len(slots) >= MAX_SLOTS
        st.button(
            "+ Adaugă întrebare",
            disabled=add_disabled,
            use_container_width=True,
            on_click=_on_add_question,
            args=(username,),
        )
        if add_disabled:
            st.caption(f"Limită maximă: {MAX_SLOTS} întrebări per imagine.")

        errors = _validate_all(slots)
        if errors:
            with st.expander("⚠️ Erori de validare", expanded=True):
                for e in errors:
                    st.error(e)

        st.button(
            "Trimite imaginea",
            type="primary",
            disabled=bool(errors),
            use_container_width=True,
            on_click=_on_submit_image,
            args=(username,),
        )


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

def _is_admin(username: str | None) -> bool:
    return bool(username and username in ADMIN_USERNAMES)


def _levenshtein(a: str, b: str) -> int:
    if a == b:
        return 0
    if not a:
        return len(b)
    if not b:
        return len(a)
    prev = list(range(len(b) + 1))
    curr = [0] * (len(b) + 1)
    for i, ca in enumerate(a, 1):
        curr[0] = i
        for j, cb in enumerate(b, 1):
            cost = 0 if ca == cb else 1
            curr[j] = min(curr[j - 1] + 1, prev[j] + 1, prev[j - 1] + cost)
        prev, curr = curr, prev
    return prev[len(b)]


def _read_all_annotations() -> list[dict]:
    import json as _json
    if not storage.ANNOTATIONS_PATH.exists():
        return []
    out: list[dict] = []
    with open(storage.ANNOTATIONS_PATH, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            out.append(_json.loads(line))
    return out


def _compute_red_flags(records: list[dict]) -> list[dict]:
    from collections import Counter
    flags: list[dict] = []

    for r in records:
        q = (r.get("question_ro") or "").strip()
        a = (r.get("answer_ro") or "").strip()
        qa_type = r.get("qa_type")
        image_id = r.get("image_id", "?")
        ql = q.lower()
        al = a.lower()

        if len(q) < 10:
            flags.append({"type": "short_question", "image_id": image_id,
                          "question": q, "detail": f"Doar {len(q)} caractere"})
        if qa_type == "open" and len(a) < 10:
            flags.append({"type": "short_answer", "image_id": image_id,
                          "question": q, "detail": f"Răspuns scurt: '{a}'"})

        if al and len(al) >= 4 and al in ql:
            flags.append({"type": "answer_in_question", "image_id": image_id,
                          "question": q, "detail": f"Răspunsul '{a}' apare în întrebare"})

        if qa_type == "mcq":
            opts = r.get("options_ro") or []
            for i, opt in enumerate(opts):
                ol = (opt or "").strip().lower()
                if not ol or ol == al:
                    continue
                if _levenshtein(ol, al) < 2:
                    flags.append({"type": "near_duplicate_distractor", "image_id": image_id,
                                  "question": q,
                                  "detail": f"Opțiunea {i + 1} ('{opt}') ≈ răspuns ('{a}')"})

    q_counts = Counter((r.get("question_ro") or "").strip().lower() for r in records)
    for q_text, count in q_counts.items():
        if count > 1 and q_text:
            flags.append({"type": "duplicate_question", "image_id": "(multiple)",
                          "question": q_text, "detail": f"Apare de {count} ori"})

    return flags


def _analytics_view() -> None:
    from collections import defaultdict
    st.title("Analytics — Romanian Cultural VLM Benchmark")

    records = _read_all_annotations()
    pool_size = len(storage.load_pool())

    with st.sidebar:
        st.markdown(f"**Admin:** `{st.session_state.username}`")
        st.divider()
        if st.button("Înapoi la adnotare", use_container_width=True):
            st.query_params.clear()
            st.rerun()
        if st.button("Ieșire", use_container_width=True):
            for k in ("username", "current_image", "qa_slots"):
                st.session_state[k] = None
            st.query_params.clear()
            st.rerun()

    if not records:
        st.info("Nu există încă adnotări trimise.")
        return

    # ---------- Progress ----------
    st.header("Progres")
    images_done = {r["image_id"] for r in records}
    mcq = sum(1 for r in records if r.get("qa_type") == "mcq")
    open_q = sum(1 for r in records if r.get("qa_type") == "open")
    last_ts = max((r.get("updated_at", "") for r in records), default="—")

    c1, c2, c3, c4 = st.columns(4)
    c1.metric("Imagini adnotate", f"{len(images_done)} / {pool_size}",
              f"{100 * len(images_done) / pool_size:.1f}%")
    c2.metric("Total Q/A", len(records))
    c3.metric("MCQ", mcq)
    c4.metric("Deschise", open_q)
    st.caption(f"Ultima trimitere: {last_ts}")

    per_user: dict[str, dict] = defaultdict(
        lambda: {"images": set(), "qa": 0, "mcq": 0, "open": 0}
    )
    for r in records:
        u = r.get("author_username", "?")
        per_user[u]["images"].add(r["image_id"])
        per_user[u]["qa"] += 1
        if r.get("qa_type") == "mcq":
            per_user[u]["mcq"] += 1
        else:
            per_user[u]["open"] += 1

    leaderboard = [
        {"Utilizator": u, "Imagini": len(d["images"]),
         "Q/A": d["qa"], "MCQ": d["mcq"], "Deschise": d["open"]}
        for u, d in sorted(per_user.items(), key=lambda kv: -kv[1]["qa"])
    ]
    st.subheader("Per utilizator")
    st.dataframe(leaderboard, use_container_width=True, hide_index=True)

    # ---------- Coverage per concept category ----------
    from collections import Counter
    pool = storage.load_pool()
    filename_to_bucket = {img["filename"]: storage.classify_image_bucket(img) for img in pool}
    pool_buckets = Counter(filename_to_bucket.values())
    done_buckets = Counter(
        bucket for fn, bucket in filename_to_bucket.items() if fn in images_done
    )
    mcq_per_bucket: Counter[str] = Counter()
    open_per_bucket: Counter[str] = Counter()
    for r in records:
        b = filename_to_bucket.get(r["image_id"], "other")
        if r.get("qa_type") == "mcq":
            mcq_per_bucket[b] += 1
        elif r.get("qa_type") == "open":
            open_per_bucket[b] += 1

    coverage = [
        {
            "Categorie": bucket,
            "În pool": pool_buckets[bucket],
            "Adnotate": done_buckets.get(bucket, 0),
            "Acoperire": (
                f"{100 * done_buckets.get(bucket, 0) / pool_buckets[bucket]:.0f}%"
                if pool_buckets[bucket] else "—"
            ),
            "MCQ": mcq_per_bucket.get(bucket, 0),
            "Deschise": open_per_bucket.get(bucket, 0),
        }
        for bucket, _ in sorted(pool_buckets.items(), key=lambda kv: -kv[1])
    ]
    st.subheader("Acoperire pe categorie")
    st.dataframe(coverage, use_container_width=True, hide_index=True)
    st.caption("Categoriile vin direct din `concept_categories` per imagine. "
               "Pentru imagini cu mai multe categorii, alegerea e deterministă pe "
               "filename. `other` = imagini fără concept_categories asociate. "
               "MCQ și Deschise numără Q/A trimise, nu imagini.")

    # ---------- Red flags ----------
    st.header("Semnale de calitate")
    flags = _compute_red_flags(records)
    if not flags:
        st.success("Niciun semnal roșu detectat.")
    else:
        from collections import Counter
        type_counts = Counter(f["type"] for f in flags)
        st.caption(" · ".join(f"`{t}`: {n}" for t, n in type_counts.most_common()))
        flag_table = [
            {"Tip": f["type"], "Imagine": f["image_id"],
             "Întrebare": (f["question"][:80] + "…") if len(f["question"]) > 80 else f["question"],
             "Detaliu": f["detail"]}
            for f in flags
        ]
        st.dataframe(flag_table, use_container_width=True, hide_index=True)


def main() -> None:
    st.set_page_config(page_title="RO Cultural VLM — Annotation", layout="wide")
    st.markdown(STICKY_CSS, unsafe_allow_html=True)
    _init_state()
    if not st.session_state.username:
        _login_view()
        return
    view = st.query_params.get("view", "annotation")
    if view == "analytics" and _is_admin(st.session_state.username):
        _analytics_view()
    else:
        _annotation_view()


if __name__ == "__main__":
    main()