"""SLM Consortium Polls — Hugging Face Space.

Astro static frontend (frontend/dist, React + React Bits components)
served at `/` by this Python script. Gradio backend mounted at
`/gradio` provides the poll API *and* Hugging Face OAuth — the Astro
app talks to it via @gradio/client on the same origin, so the session
cookie carries HF auth and every mutating call is org-gated.

Auth model:
- Space metadata sets `hf_oauth: true` (see README.md).
- `check_membership()` verifies `slmconsortium` membership via
  `whoami(user_token)["orgs"]`, with a public
  `list_organization_members` fallback.
- Signed-out users can read polls/results; only members can vote/create.
"""

from __future__ import annotations

import json
import os
import threading
import uuid
from datetime import datetime, timezone
from typing import Dict, List, Optional, Tuple

import gradio as gr
from fastapi import FastAPI, Request
from fastapi.responses import RedirectResponse
from fastapi.staticfiles import StaticFiles
from huggingface_hub import HfApi, whoami

REQUIRED_ORG = os.getenv("REQUIRED_ORG", "slmconsortium")
MEMBERSHIP_CACHE_TTL = 300  # seconds

_HERE = os.path.dirname(os.path.abspath(__file__))
# Persistent storage on Spaces lives under /data (when enabled).
# Fall back to the app directory for local dev / ephemeral disk.
_DATA_DIR = "/data" if os.path.isdir("/data") else _HERE
POLLS_FILE = os.path.join(_DATA_DIR, "polls.json")
DIST_DIR = os.path.join(_HERE, "frontend", "dist")
_lock = threading.Lock()
_member_cache: Dict[str, Tuple[bool, float]] = {}


# ---------------------------------------------------------------- storage

def load_polls() -> Dict:
    if not os.path.exists(POLLS_FILE):
        return {}
    try:
        with open(POLLS_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
            return data if isinstance(data, dict) else {}
    except (json.JSONDecodeError, OSError):
        return {}


def save_polls(polls: Dict) -> None:
    tmp = POLLS_FILE + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(polls, f, indent=2, ensure_ascii=False)
    os.replace(tmp, POLLS_FILE)


def total_votes(poll: Dict) -> int:
    return len(poll.get("voters", {}))


def count_votes(poll: Dict) -> List[int]:
    n = len(poll.get("options", []))
    counts = [0] * n
    for idx in poll.get("voters", {}).values():
        if isinstance(idx, int) and 0 <= idx < n:
            counts[idx] += 1
    return counts


def parse_options(raw: str) -> List[str]:
    options = [o.strip() for o in (raw or "").splitlines() if o.strip()]
    if len(options) <= 1 and raw and "," in raw:
        options = [o.strip() for o in raw.split(",") if o.strip()]
    return options


def poll_to_dict(poll: Dict, username: Optional[str]) -> Dict:
    counts = count_votes(poll)
    n = len(poll.get("options", []))
    return {
        "id": poll["id"],
        "question": poll["question"],
        "options": poll["options"],
        "counts": counts,
        "total": sum(counts),
        "created_by": poll.get("created_by", "?"),
        "created_at": poll.get("created_at", ""),
        "my_vote": (poll.get("voters", {}).get(username) if username else None),
        # username -> option index (sanitized)
        "voters": {
            u: idx
            for u, idx in poll.get("voters", {}).items()
            if isinstance(idx, int) and 0 <= idx < n
        },
    }


# ------------------------------------------------------------------ auth

def check_membership(
    profile: gr.OAuthProfile | None,
    token: gr.OAuthToken | None,
) -> Tuple[bool, str, str]:
    """Return (allowed, username, status_message)."""
    if profile is None:
        return False, "", "🔒 Please **sign in with Hugging Face** to create polls or vote."

    username = profile.username or ""
    now = datetime.now(timezone.utc).timestamp()
    cached = _member_cache.get(username.lower())
    if cached is not None:
        allowed, ts = cached
        if now - ts < MEMBERSHIP_CACHE_TTL:
            msg = (
                f"✅ Signed in as **@{username}** — member of `{REQUIRED_ORG}`."
                if allowed
                else f"⛔ Signed in as **@{username}** — not a member of `{REQUIRED_ORG}`. "
                     "Creating polls and voting are disabled."
            )
            return allowed, username, msg

    # 1) Preferred: whoami with the user's own OAuth token.
    #    Sees private org memberships too.
    if token is not None and getattr(token, "token", None):
        try:
            info = whoami(token.token)
            orgs = info.get("orgs", []) or []
            names = [
                (o.get("name", "") if isinstance(o, dict) else str(o)).lower()
                for o in orgs
            ]
            allowed = REQUIRED_ORG.lower() in names
            _member_cache[username.lower()] = (allowed, now)
            if allowed:
                return True, username, f"✅ Signed in as **@{username}** — member of `{REQUIRED_ORG}`."
            return False, username, (
                f"⛔ Signed in as **@{username}** — not a member of `{REQUIRED_ORG}`. "
                "Creating polls and voting are disabled."
            )
        except Exception:
            pass  # fall through to public check

    # 2) Fallback: public org member list (works without a user token
    #    for public orgs; also covers local-dev where token may be None).
    try:
        api = HfApi()
        members = [m.username.lower() for m in api.list_organization_members(REQUIRED_ORG)]
        allowed = username.lower() in members
        _member_cache[username.lower()] = (allowed, now)
        if allowed:
            return True, username, f"✅ Signed in as **@{username}** — member of `{REQUIRED_ORG}`."
        return False, username, (
            f"⛔ Signed in as **@{username}** — not a member of `{REQUIRED_ORG}`. "
            "Creating polls and voting are disabled."
        )
    except Exception as e:
        return False, username, (
            f"⚠️ Signed in as **@{username}**, but membership in `{REQUIRED_ORG}` "
            f"could not be verified ({e}). Try again shortly."
        )


# ------------------------------------------------- JSON API (Astro app)

def api_me(
    profile: gr.OAuthProfile | None,
    token: gr.OAuthToken | None,
) -> Dict:
    allowed, username, _ = check_membership(profile, token)
    return {"username": username or None, "member": allowed}


def api_polls(
    profile: gr.OAuthProfile | None,
    token: gr.OAuthToken | None,
) -> Dict:
    allowed, username, _ = check_membership(profile, token) if profile else (False, "", "")
    polls = load_polls()
    items = sorted(polls.values(), key=lambda p: p.get("created_at", ""), reverse=True)
    return {
        "me": {"username": username or None, "member": allowed},
        "polls": [poll_to_dict(p, username or None) for p in items],
    }


def api_vote(
    poll_id: str,
    choice_idx: float,
    profile: gr.OAuthProfile | None,
    token: gr.OAuthToken | None,
) -> Dict:
    allowed, username, status = check_membership(profile, token)
    polls = load_polls()
    poll = polls.get(poll_id or "")
    if not allowed:
        return {"ok": False, "message": status, "poll": poll_to_dict(poll, None) if poll else None}
    if poll is None:
        return {"ok": False, "message": "That poll no longer exists — refresh the list.", "poll": None}
    try:
        idx = int(choice_idx)
    except (TypeError, ValueError):
        idx = -1
    if not 0 <= idx < len(poll["options"]):
        return {"ok": False, "message": "Invalid option.", "poll": poll_to_dict(poll, username)}

    with _lock:
        polls = load_polls()
        poll = polls.get(poll_id)
        if poll is None:
            return {"ok": False, "message": "That poll no longer exists — refresh the list.", "poll": None}
        prev = poll.setdefault("voters", {}).get(username)
        poll["voters"][username] = idx
        save_polls(polls)

    msg = f"Vote counted for “{poll['options'][idx]}” as @{username}."
    if prev is not None and prev != idx and 0 <= prev < len(poll["options"]):
        msg += f" (changed from “{poll['options'][prev]}”)"
    return {"ok": True, "message": msg, "poll": poll_to_dict(polls[poll_id], username)}


def _validate_poll(question: str, options_raw: str) -> Tuple[str, List[str], Optional[str]]:
    """Return (question, options, error_message_or_None)."""
    question = (question or "").strip()
    options = parse_options(options_raw or "")
    if len(question) < 3:
        return question, options, "Give your poll a question (min 3 characters)."
    if len(options) < 2:
        return question, options, "Provide at least 2 options (one per line)."
    if len(options) > 20:
        return question, options, "Max 20 options per poll."
    if len(set(o.lower() for o in options)) != len(options):
        return question, options, "Options must be unique."
    return question, options, None


def api_create(
    question: str,
    options_raw: str,
    profile: gr.OAuthProfile | None,
    token: gr.OAuthToken | None,
) -> Dict:
    allowed, username, status = check_membership(profile, token)
    if not allowed:
        return {"ok": False, "message": status, "poll": None}

    question, options, err = _validate_poll(question, options_raw)
    if err:
        return {"ok": False, "message": err, "poll": None}

    with _lock:
        polls = load_polls()
        pid = uuid.uuid4().hex[:8]
        polls[pid] = {
            "id": pid,
            "question": question,
            "options": options,
            "created_by": username,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "voters": {},
        }
        save_polls(polls)
    return {
        "ok": True,
        "message": f"Poll created by @{username}!",
        "poll": poll_to_dict(polls[pid], username),
    }


def api_edit(
    poll_id: str,
    question: str,
    options_raw: str,
    profile: gr.OAuthProfile | None,
    token: gr.OAuthToken | None,
) -> Dict:
    allowed, username, status = check_membership(profile, token)
    if not allowed:
        return {"ok": False, "message": status, "poll": None}

    question, options, err = _validate_poll(question, options_raw)
    if err:
        return {"ok": False, "message": err, "poll": None}

    with _lock:
        polls = load_polls()
        poll = polls.get(poll_id or "")
        if poll is None:
            return {"ok": False, "message": "That poll no longer exists — refresh the list.", "poll": None}
        # Keep votes for choices whose text is unchanged; drop votes for
        # removed/edited options (indices may shift, so remap by text).
        old_opts = poll["options"]
        remap: Dict[int, int] = {}
        for old_idx, opt in enumerate(old_opts):
            if opt in options:
                remap[old_idx] = options.index(opt)
        votes_reset = old_opts != options
        kept = 0
        for u, idx in list(poll.get("voters", {}).items()):
            new_idx = remap.get(idx)
            if new_idx is None:
                poll["voters"].pop(u, None)
            else:
                poll["voters"][u] = new_idx
                kept += 1
        poll["question"] = question
        poll["options"] = options
        save_polls(polls)
        updated = poll_to_dict(poll, username)

    msg = f"Poll updated by @{username}."
    if votes_reset:
        msg += f" Kept {kept} vote(s) for unchanged options; votes for edited/removed options were dropped."
    return {"ok": True, "message": msg, "poll": updated}


def api_delete(
    poll_id: str,
    profile: gr.OAuthProfile | None,
    token: gr.OAuthToken | None,
) -> Dict:
    allowed, username, status = check_membership(profile, token)
    if not allowed:
        return {"ok": False, "message": status}

    with _lock:
        polls = load_polls()
        poll = polls.pop(poll_id or "", None)
        if poll is None:
            return {"ok": False, "message": "That poll no longer exists — refresh the list."}
        save_polls(polls)
    return {"ok": True, "message": f"Poll “{poll['question']}” deleted by @{username}."}


# ------------------------------------------------------- classic UI views

def poll_label(poll: Dict) -> str:
    return f"{poll['question']}  —  {total_votes(poll)} vote(s)  [{poll['id'][:6]}]"


def dropdown_choices(polls: Dict) -> List[Tuple[str, str]]:
    items = sorted(polls.values(), key=lambda p: p.get("created_at", ""), reverse=True)
    return [(poll_label(p), p["id"]) for p in items]


def format_results(poll: Optional[Dict]) -> str:
    if poll is None:
        return "Select a poll to see live results."
    counts = count_votes(poll)
    total = sum(counts)
    lines = [f"### 📊 {poll['question']}", ""]
    if total == 0:
        lines.append("_No votes yet — be the first!_")
    by_opt: Dict[int, List[str]] = {}
    for u, idx in poll.get("voters", {}).items():
        if isinstance(idx, int) and 0 <= idx < len(poll["options"]):
            by_opt.setdefault(idx, []).append(u)
    for i, (opt, c) in enumerate(zip(poll["options"], counts)):
        pct = (c / total * 100) if total else 0
        bar = "█" * int(round(pct / 5)) + "░" * (20 - int(round(pct / 5)))
        lines.append(f"- **{opt}** — {c} vote(s) ({pct:.1f}%)  `{bar}`")
        voters = sorted(by_opt.get(i, []))
        if voters:
            lines.append(f"  - Voters: " + ", ".join(f"**@{u}**" for u in voters))
    lines.append("")
    lines.append(f"Total votes: **{total}** · Created by **@{poll.get('created_by', '?')}**")
    return "\n".join(lines)


def refresh_polls() -> Tuple[gr.Dropdown, str]:
    polls = load_polls()
    choices = dropdown_choices(polls)
    dd = gr.Dropdown(choices=choices, value=choices[0][1] if choices else None)
    if not choices:
        return dd, "No polls yet. Members of `slmconsortium` can create one in the **Create poll** tab."
    poll = polls.get(choices[0][1])
    return dd, format_results(poll)


def on_select_poll(
    poll_id: Optional[str],
) -> Tuple[gr.Radio, str]:
    polls = load_polls()
    poll = polls.get(poll_id) if poll_id else None
    if poll is None:
        return gr.Radio(choices=[], value=None), "Select a poll to see live results."
    radio = gr.Radio(choices=poll["options"], value=None)
    return radio, format_results(poll)


def on_vote(
    poll_id: Optional[str],
    choice: Optional[str],
    profile: gr.OAuthProfile | None,
    token: gr.OAuthToken | None,
) -> Tuple[str, str]:
    allowed, username, status = check_membership(profile, token)
    if not allowed:
        return status, format_results(load_polls().get(poll_id) if poll_id else None)
    if not poll_id:
        return "⚠️ Select a poll first.", "Select a poll to see live results."
    if not choice:
        return "⚠️ Pick an option before voting.", format_results(load_polls().get(poll_id))

    with _lock:
        polls = load_polls()
        poll = polls.get(poll_id)
        if poll is None:
            return "⚠️ That poll no longer exists. Hit Refresh.", "Select a poll to see live results."
        try:
            idx = poll["options"].index(choice)
        except ValueError:
            return "⚠️ Invalid option.", format_results(poll)
        prev = poll.setdefault("voters", {}).get(username)
        poll["voters"][username] = idx
        save_polls(polls)

    msg = f"✅ Vote counted for **{choice}** as **@{username}**."
    if prev is not None and prev != idx:
        msg += f" (changed from **{poll['options'][prev]}**)"
    elif prev == idx:
        msg += " (unchanged)"
    polls = load_polls()
    return msg, format_results(polls.get(poll_id))


def on_create(
    question: str,
    options_raw: str,
    profile: gr.OAuthProfile | None,
    token: gr.OAuthToken | None,
) -> Tuple[str, gr.Dropdown, gr.Radio, str, str, str]:
    allowed, username, status = check_membership(profile, token)
    polls = load_polls()
    choices = dropdown_choices(polls)
    dd = gr.Dropdown(choices=choices, value=choices[0][1] if choices else None)
    current = polls.get(choices[0][1]) if choices else None
    radio = gr.Radio(choices=current["options"] if current else [])
    results = format_results(current)

    if not allowed:
        return status, dd, radio, results, question, options_raw

    question = (question or "").strip()
    options = parse_options(options_raw or "")

    if len(question) < 3:
        return "⚠️ Give your poll a question (min 3 characters).", dd, radio, results, question, options_raw
    if len(options) < 2:
        return "⚠️ Provide at least **2 options** (one per line).", dd, radio, results, question, options_raw
    if len(options) > 20:
        return "⚠️ Max 20 options per poll.", dd, radio, results, question, options_raw
    if len(set(o.lower() for o in options)) != len(options):
        return "⚠️ Options must be unique.", dd, radio, results, question, options_raw

    with _lock:
        polls = load_polls()
        pid = uuid.uuid4().hex[:8]
        polls[pid] = {
            "id": pid,
            "question": question,
            "options": options,
            "created_by": username,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "voters": {},
        }
        save_polls(polls)
        choices = dropdown_choices(polls)

    dd = gr.Dropdown(choices=choices, value=pid)
    radio = gr.Radio(choices=options, value=None)
    results = format_results(polls[pid])
    return f"✅ Poll created by **@{username}**! Share it with fellow `{REQUIRED_ORG}` members.", dd, radio, results, "", ""


def on_load(
    profile: gr.OAuthProfile | None,
    token: gr.OAuthToken | None,
) -> Tuple[str, gr.Dropdown, gr.Radio, str]:
    _, _, status = check_membership(profile, token) if profile else (False, "", "🔒 Please **sign in with Hugging Face** to create polls or vote.")
    polls = load_polls()
    choices = dropdown_choices(polls)
    dd = gr.Dropdown(choices=choices, value=choices[0][1] if choices else None)
    current = polls.get(choices[0][1]) if choices else None
    radio = gr.Radio(choices=current["options"] if current else [])
    return status, dd, radio, format_results(current)


# ------------------------------------------------------------------ blocks

# Link for reaching the Astro frontend from the gradio UI. On a private
# Space a bare "/" (hf.space root) 404s in a fresh tab — the HF auth
# handshake only happens via the huggingface.co Space page.
_space_id = os.getenv("SPACE_ID", "")
_astro_link = f"https://huggingface.co/spaces/{_space_id}" if _space_id else "/"

with gr.Blocks(title="SLM Consortium Polls") as demo:
    gr.Markdown(
        "# 🗳️ SLM Consortium Polls\n"
        "Create polls and vote. **Sign-in with Hugging Face is required**, and only "
        f"members of the `{REQUIRED_ORG}` organization can create polls or vote. "
        "Everyone can view live results.\n\n"
        f"Prefer the new look? Use the [Astro frontend]({_astro_link}) at the Space root."
    )
    with gr.Row():
        gr.LoginButton(min_width=50)
        status_md = gr.Markdown("🔒 Please **sign in with Hugging Face** to create polls or vote.")

    with gr.Tab("🗳️ Vote"):
        poll_dropdown = gr.Dropdown(label="Choose a poll", choices=[], interactive=True)
        option_radio = gr.Radio(label="Your choice", choices=[])
        with gr.Row():
            vote_btn = gr.Button("Vote", variant="primary")
            refresh_btn = gr.Button("🔄 Refresh")
        vote_msg = gr.Markdown()
        results_md = gr.Markdown("Select a poll to see live results.")

    with gr.Tab("➕ Create poll"):
        gr.Markdown(
            f"Only `{REQUIRED_ORG}` members can create polls. "
            "Put each option on its own line (or comma-separated)."
        )
        q_input = gr.Textbox(label="Question", placeholder="e.g. Which meeting time works best?")
        o_input = gr.Textbox(
            label="Options (one per line)",
            lines=4,
            placeholder="Monday 10:00 UTC\nTuesday 14:00 UTC\nFriday 09:00 UTC",
        )
        create_btn = gr.Button("Create poll", variant="primary")
        create_msg = gr.Markdown()

    # Wiring (classic UI)
    demo.load(on_load, inputs=None, outputs=[status_md, poll_dropdown, option_radio, results_md])
    poll_dropdown.change(on_select_poll, inputs=poll_dropdown, outputs=[option_radio, results_md])
    refresh_btn.click(refresh_polls, inputs=None, outputs=[poll_dropdown, results_md])
    vote_btn.click(on_vote, inputs=[poll_dropdown, option_radio], outputs=[vote_msg, results_md])
    create_btn.click(
        on_create,
        inputs=[q_input, o_input],
        outputs=[create_msg, poll_dropdown, option_radio, results_md, q_input, o_input],
    )

    # Hidden JSON API for the Astro frontend (same-origin, same HF session).
    api_me_btn = gr.Button(visible=False)
    api_me_out = gr.JSON(visible=False)
    api_me_btn.click(api_me, inputs=None, outputs=api_me_out, api_name="me")

    api_polls_btn = gr.Button(visible=False)
    api_polls_out = gr.JSON(visible=False)
    api_polls_btn.click(api_polls, inputs=None, outputs=api_polls_out, api_name="polls")

    api_vote_poll_id = gr.Textbox(visible=False)
    api_vote_choice = gr.Number(visible=False, precision=0)
    api_vote_out = gr.JSON(visible=False)
    api_vote_btn = gr.Button(visible=False)
    api_vote_btn.click(
        api_vote,
        inputs=[api_vote_poll_id, api_vote_choice],
        outputs=api_vote_out,
        api_name="vote",
    )

    api_create_q = gr.Textbox(visible=False)
    api_create_opts = gr.Textbox(visible=False)
    api_create_out = gr.JSON(visible=False)
    api_create_btn = gr.Button(visible=False)
    api_create_btn.click(
        api_create,
        inputs=[api_create_q, api_create_opts],
        outputs=api_create_out,
        api_name="create_poll",
    )

    api_edit_pid = gr.Textbox(visible=False)
    api_edit_q = gr.Textbox(visible=False)
    api_edit_opts = gr.Textbox(visible=False)
    api_edit_out = gr.JSON(visible=False)
    api_edit_btn = gr.Button(visible=False)
    api_edit_btn.click(
        api_edit,
        inputs=[api_edit_pid, api_edit_q, api_edit_opts],
        outputs=api_edit_out,
        api_name="edit_poll",
    )

    api_del_pid = gr.Textbox(visible=False)
    api_del_out = gr.JSON(visible=False)
    api_del_btn = gr.Button(visible=False)
    api_del_btn.click(
        api_delete,
        inputs=[api_del_pid],
        outputs=api_del_out,
        api_name="delete_poll",
    )


# ------------------------------------------------- app: static + gradio

fastapi_app = FastAPI(title="SLM Consortium Polls")

# Gradio 6 mounts don't redirect /gradio -> /gradio/, so the bare path 404s.
# Add the redirect before the mount takes over the prefix.
@fastapi_app.get("/gradio")
def gradio_redirect() -> RedirectResponse:
    return RedirectResponse(url="/gradio/", status_code=307)


# Gradio's OAuth routes are hardcoded at the app root ("/login/huggingface",
# "/login/callback", "/logout") and its Spaces flow forces the callback to
# https://<space_host>/login/callback — outside the /gradio mount. Forward
# those bare paths into the mounted app so HF OAuth works.
def _forward_to_gradio(path: str):
    async def _forward(request: Request) -> RedirectResponse:
        qs = "&".join(
            f"{k}={v}" for k, v in request.query_params.multi_items()
        )
        url = f"/gradio{path}" + (f"?{qs}" if qs else "")
        return RedirectResponse(url, status_code=307)

    return _forward


fastapi_app.get("/login/huggingface")(_forward_to_gradio("/login/huggingface"))
fastapi_app.get("/login/callback")(_forward_to_gradio("/login/callback"))
fastapi_app.get("/logout")(_forward_to_gradio("/logout"))


# root_path makes gradio generate internal links (OAuth login, etc.)
# under /gradio instead of the app root. ssr_mode=False is required on
# Spaces: SSR's node server would grab port 7860 and starve uvicorn.
fastapi_app = gr.mount_gradio_app(
    fastapi_app,
    demo,
    path="/gradio",
    root_path="/gradio",
    ssr_mode=False,
)


@fastapi_app.get("/healthz")
def healthz() -> Dict[str, str]:
    return {
        "status": "ok",
        "data_dir": _DATA_DIR,
        "polls_file": POLLS_FILE,
        "polls_on_disk": str(os.path.exists(POLLS_FILE)),
    }


if os.path.isdir(DIST_DIR):
    fastapi_app.mount("/", StaticFiles(directory=DIST_DIR, html=True), name="frontend")
else:
    @fastapi_app.get("/")
    def missing_build() -> Dict[str, str]:
        return {
            "status": "frontend not built",
            "hint": "Run `npm run build` in frontend/ (output committed to frontend/dist/).",
            "gradio": "/gradio",
        }


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(fastapi_app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")))