#!/usr/bin/env python3
"""
Human evaluation web app for Hugging Face Spaces.

Shows stitched pair images in random order per session and logs Yes/No answers
with response time per image (JSON key = image filename).
"""

from __future__ import annotations

import json
import os

# HF Docker: keep localhost checks out of HTTP(S)_PROXY (Gradio launch self-test).
_NO_PROXY = "localhost,127.0.0.1,::1,0.0.0.0"
os.environ.setdefault("NO_PROXY", _NO_PROXY)
os.environ.setdefault("no_proxy", _NO_PROXY)
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
import random
import re
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import gradio as gr

APP_DIR = Path(__file__).resolve().parent
HUMAN_EVAL_DIR = Path(os.environ.get("HUMAN_EVAL_DIR", APP_DIR / "human_eval"))
RESPONSES_DIR = Path(os.environ.get("RESPONSES_DIR", "/data/responses"))
if not RESPONSES_DIR.parent.exists():
    RESPONSES_DIR = APP_DIR / "responses"

STUDY_TITLE = "Visual Transformation Study: Image Rotation"
QUESTION_TEXT = "If I rotate the first image, can I get the second image?"
THANK_YOU_TEXT = "THANK YOU!!!"
CHOICES = ["Yes", "No"]


def _intro_markdown(stimuli_count: int) -> str:
    return f"""# {STUDY_TITLE}

You will be shown {stimuli_count} pairs of images. For each pair, you need to determine if the second image (Right image) is simply a rotated version of the first image (Left image).

**Question:** "{QUESTION_TEXT}"

Answer with **"Yes"** or **"No"**, then click **Next**.

Use the **same username** to come back and resume where you left off.
"""


def _question_markdown() -> str:
    return f'**Question:** "{QUESTION_TEXT}"'


def _enable_next(answer: Optional[str]):
    return gr.update(interactive=answer in CHOICES)


def _utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _safe_username(name: str) -> str:
    slug = re.sub(r"[^\w.\-]+", "_", name.strip())
    return slug[:64] or "anonymous"


def discover_stimuli(root: Path) -> List[Dict[str, str]]:
    if not root.is_dir():
        raise FileNotFoundError(
            f"human_eval folder not found: {root}\n"
            "Copy your stimuli into human_eval_hf_space/human_eval/ before deploy."
        )
    items: List[Dict[str, str]] = []
    for png in sorted(root.glob("Q*/*_pair.png")):
        items.append(
            {
                "filename": png.name,
                "path": str(png.resolve()),
                "folder": png.parent.name,
            }
        )
    if not items:
        raise FileNotFoundError(f"No *_pair.png files under {root}/Q*/")
    return items


def _response_path(username: str) -> Path:
    RESPONSES_DIR.mkdir(parents=True, exist_ok=True)
    return RESPONSES_DIR / f"{_safe_username(username)}.json"


def _load_response_file(username: str) -> Dict[str, Any]:
    path = _response_path(username)
    if path.is_file():
        with open(path, encoding="utf-8") as f:
            return json.load(f)
    return {
        "username": _safe_username(username),
        "created_at": _utc_now(),
        "updated_at": None,
        "responses": {},
    }


def _save_response_file(username: str, data: Dict[str, Any]) -> str:
    data["username"] = _safe_username(username)
    data["updated_at"] = _utc_now()
    path = _response_path(username)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)
    return str(path)


def _build_order(username: str, all_items: List[Dict[str, str]]) -> List[Dict[str, str]]:
    """Return question order for this user, reusing a saved shuffle when resuming."""
    record = _load_response_file(username)
    by_name = {item["filename"]: item for item in all_items}
    current_names = set(by_name)

    saved_names = record.get("order")
    if saved_names:
        order: List[Dict[str, str]] = []
        seen: set[str] = set()
        for filename in saved_names:
            if filename in by_name:
                order.append(by_name[filename])
                seen.add(filename)
        new_names = sorted(current_names - seen)
        if new_names:
            rng = random.Random(username)
            extra = [by_name[name] for name in new_names]
            rng.shuffle(extra)
            order.extend(extra)
        if {item["filename"] for item in order} == current_names:
            if new_names:
                record["order"] = [item["filename"] for item in order]
                _save_response_file(username, record)
            return order

    rng = random.Random(username)
    order = list(all_items)
    rng.shuffle(order)
    record["order"] = [item["filename"] for item in order]
    if not record.get("created_at"):
        record["created_at"] = _utc_now()
    _save_response_file(username, record)
    return order


def _resume_index(order: List[Dict[str, str]], responses: Dict[str, Any]) -> int:
    for i, item in enumerate(order):
        if item["filename"] not in responses:
            return i
    return len(order)


def start_session(username: str) -> Tuple[Any, ...]:
    if not username or not username.strip():
        raise gr.Error("Please enter a username before starting.")

    name = username.strip()
    all_items = discover_stimuli(HUMAN_EVAL_DIR)
    order = _build_order(name, all_items)
    record = _load_response_file(name)
    idx = _resume_index(order, record.get("responses", {}))

    status_msg = ""
    answered = idx
    total = len(order)
    if idx >= total:
        status_msg = f"You already completed all {total} questions."
    elif answered > 0:
        status_msg = f"Resuming — **{answered}** of **{total}** already answered."

    state = {
        "username": name,
        "order": order,
        "index": idx,
        "question_started_at": time.perf_counter(),
        "session_started_at": record.get("created_at") or _utc_now(),
    }
    if idx > 0 or record.get("instructions_acknowledged_at"):
        return _render_question(state, status_msg=status_msg)
    return _render_consent(state)


def _render_consent(state: Dict[str, Any]) -> Tuple[Any, ...]:
    return (
        state,
        gr.update(visible=False),
        gr.update(visible=True),
        gr.update(visible=False),
        gr.update(visible=False),
        gr.update(value=state["username"]),
        gr.update(value=None),
        gr.update(value=_question_markdown()),
        gr.update(value=None, choices=CHOICES),
        gr.update(value=""),
        gr.update(value=""),
        gr.update(visible=False, interactive=False),
    )


def acknowledge_instructions(state: Optional[Dict[str, Any]]) -> Tuple[Any, ...]:
    if state is None:
        raise gr.Error("Enter a username and click **Start / Resume** first.")

    record = _load_response_file(state["username"])
    record["instructions_acknowledged_at"] = _utc_now()
    _save_response_file(state["username"], record)
    return _render_question(state)


def _render_question(
    state: Optional[Dict[str, Any]],
    status_msg: str = "",
) -> Tuple[Any, ...]:
    if state is None:
        return (
            None,
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(value=""),
            gr.update(value=None),
            gr.update(value=_question_markdown()),
            gr.update(value=None, choices=CHOICES),
            gr.update(value=""),
            gr.update(value="Enter a username and click **Start / Resume**."),
            gr.update(visible=False, interactive=False),
        )

    total = len(state["order"])
    idx = state["index"]
    if idx >= total:
        return (
            state,
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(value=state["username"]),
            gr.update(value=None),
            gr.update(value=f"### {THANK_YOU_TEXT}"),
            gr.update(value=None, choices=CHOICES),
            gr.update(value=""),
            gr.update(
                value=status_msg
                or f"You completed all **{total}** questions. **{THANK_YOU_TEXT}**"
            ),
            gr.update(visible=False, interactive=False),
        )

    item = state["order"][idx]
    progress_text = f"Question **{idx + 1} / {total}**"
    state["question_started_at"] = time.perf_counter()

    return (
        state,
        gr.update(visible=False),
        gr.update(visible=False),
        gr.update(visible=True),
        gr.update(visible=False),
        gr.update(value=state["username"]),
        gr.update(value=item["path"]),
        gr.update(value=_question_markdown()),
        gr.update(value=None, choices=CHOICES),
        gr.update(value=progress_text),
        gr.update(value=status_msg),
        gr.update(visible=True, interactive=False),
    )


def submit_answer(
    state: Optional[Dict[str, Any]],
    answer: Optional[str],
) -> Tuple[Any, ...]:
    if state is None:
        raise gr.Error("Click **Start** before answering.")

    total = len(state["order"])
    idx = state["index"]
    if idx >= total:
        return _render_question(state)

    if answer not in CHOICES:
        raise gr.Error("Please select Yes or No.")

    item = state["order"][idx]
    elapsed = time.perf_counter() - state["question_started_at"]
    filename = item["filename"]

    record = _load_response_file(state["username"])
    record["responses"][filename] = {
        "answer": answer,
        "time_seconds": round(elapsed, 3),
        "question_index": idx + 1,
        "question_folder": item["folder"],
        "answered_at": _utc_now(),
    }
    _save_response_file(state["username"], record)

    state["index"] = idx + 1
    return _render_question(state)


def build_ui() -> gr.Blocks:
    stimuli_count = len(discover_stimuli(HUMAN_EVAL_DIR))

    with gr.Blocks(
        title=STUDY_TITLE,
        theme=gr.themes.Soft(),
    ) as demo:
        gr.Markdown(_intro_markdown(stimuli_count))

        state = gr.State(None)

        with gr.Group(visible=True) as setup_panel:
            username = gr.Textbox(
                label="Username",
                placeholder="e.g. your name or participant ID",
            )
            start_btn = gr.Button("Start / Resume", variant="primary")

        with gr.Group(visible=False) as consent_panel:
            gr.Markdown(
                "Please read the instructions above, then confirm before starting the questions."
            )
            consent_btn = gr.Button("Yes, I have read the instructions", variant="primary")

        with gr.Group(visible=False) as question_panel:
            user_display = gr.Textbox(label="Participant", interactive=False)
            image = gr.Image(type="filepath", show_label=False)
            prompt = gr.Markdown(value=_question_markdown())
            answer = gr.Radio(choices=CHOICES, label='Answer with "Yes" or "No"', value=None)
            progress = gr.Markdown(value="")
            status = gr.Markdown(value="")
            next_btn = gr.Button("Next", variant="primary", interactive=False)

        with gr.Group(visible=False) as done_panel:
            gr.Markdown(f"### {THANK_YOU_TEXT}")

        outputs = [
            state,
            setup_panel,
            consent_panel,
            question_panel,
            done_panel,
            user_display,
            image,
            prompt,
            answer,
            progress,
            status,
            next_btn,
        ]

        start_btn.click(fn=start_session, inputs=[username], outputs=outputs)
        consent_btn.click(fn=acknowledge_instructions, inputs=[state], outputs=outputs)
        answer.change(fn=_enable_next, inputs=[answer], outputs=[next_btn])
        next_btn.click(fn=submit_answer, inputs=[state, answer], outputs=outputs)

    return demo


def _server_name() -> str:
    """HF Spaces must bind 0.0.0.0; local dev should use 127.0.0.1 (0.0.0.0 often shows a blank page)."""
    if os.environ.get("SPACE_ID") or Path("/data").is_dir():
        return "0.0.0.0"
    return os.environ.get("GRADIO_SERVER_NAME", "127.0.0.1")


def _is_hf_or_docker() -> bool:
    return bool(os.environ.get("SPACE_ID")) or Path("/data").is_dir()


def _prepare_gradio_launch() -> None:
    """Docker/HF bind 0.0.0.0; Gradio's HEAD self-check to that URL often fails."""
    if not _is_hf_or_docker():
        return
    import gradio.networking as networking

    networking.url_ok = lambda _url: True


if __name__ == "__main__":
    print(f"Stimuli: {HUMAN_EVAL_DIR}")
    print(f"Responses: {RESPONSES_DIR}")
    app = build_ui()
    port = int(os.environ.get("PORT", "7860"))
    host = _server_name()
    print(f"Listening on http://{host}:{port}/")
    _prepare_gradio_launch()
    app.launch(
        server_name=host,
        server_port=port,
        share=False,
        show_error=True,
        show_api=False,
    )