"""
Image Similarity Rating App
----------------------------
Reads pairs.csv (committed to the Space repo) and shows all pairs
in random order to each user. No repetitions within a session.

pairs.csv columns:
  describer, generator, experiment, episode, turn,
  original_image_url, generated_image_url
"""

import io
import os
import uuid
import random
import pandas as pd
import gradio as gr
from datetime import datetime
from datasets import Dataset
from huggingface_hub import HfApi

# ── Config ────────────────────────────────────────────────────────────────────
HF_TOKEN        = os.environ.get("HF_TOKEN", "")
HF_DATASET_REPO = os.environ.get("HF_DATASET_REPO", "")  # where votes are saved
CSV_PATH        = "pairs.csv"                              # committed to Space repo

# ── Load pairs CSV once at startup ───────────────────────────────────────────
print("Loading pairs.csv ...")
_pairs_df = pd.read_csv(CSV_PATH)
_pairs = _pairs_df.to_dict(orient="records")
print(f"Loaded {len(_pairs)} pairs.")


# ── Persistence ───────────────────────────────────────────────────────────────
VOTES_FILE = "votes.parquet"  # single file in the results repo

def save_votes_to_hub(votes: list[dict]):
    """
    Append this session's votes to a single votes.parquet in the results repo.
    Strategy: download existing file -> concat -> upload back.
    """
    if not HF_DATASET_REPO or not HF_TOKEN:
        print("HF_DATASET_REPO or HF_TOKEN not set -- votes not saved remotely.")
        return
    try:
        api = HfApi(token=HF_TOKEN)
        new_df = pd.DataFrame(votes)

        # Try to download the existing parquet and append
        try:
            existing_path = api.hf_hub_download(
                repo_id=HF_DATASET_REPO,
                repo_type="dataset",
                filename=VOTES_FILE,
            )
            existing_df = pd.read_parquet(existing_path)
            combined_df = pd.concat([existing_df, new_df], ignore_index=True)
        except Exception:
            # File doesn't exist yet -- first run
            combined_df = new_df

        buf = io.BytesIO()
        combined_df.to_parquet(buf, index=False)
        buf.seek(0)

        api.upload_file(
            path_or_fileobj=buf,
            path_in_repo=VOTES_FILE,
            repo_id=HF_DATASET_REPO,
            repo_type="dataset",
        )
        print(f"Appended {len(votes)} votes to {HF_DATASET_REPO}/{VOTES_FILE} "
              f"(total rows: {len(combined_df)})")
    except Exception as ex:
        print(f"Failed to save votes: {ex}")


# ── CSS ───────────────────────────────────────────────────────────────────────
CSS = """
/* Light mode */
@media (prefers-color-scheme: light) {
    body, .gradio-container { background: #f5f5f5; color: #111; }
    .instructions           { background: #fff; border-color: #ddd; color: #444; }
    .instructions strong    { color: #111; }
    .scale-list li          { color: #555; }
    .done-banner            { background: #fff; border-color: #bbb; }
    .done-banner p          { color: #555; }
    .progress-track         { background: #ddd; }
    .progress-fill          { background: #333; }
}

/* Dark mode */
@media (prefers-color-scheme: dark) {
    body, .gradio-container { background: #141414; color: #e8e8e8; }
    .instructions           { background: #1e1e1e; border-color: #333; color: #bbb; }
    .instructions strong    { color: #eee; }
    .scale-list li          { color: #999; }
    .scale-val              { background: #e8e8e8 !important; color: #111 !important; }
    .done-banner            { background: #1e1e1e; border-color: #444; }
    .done-banner h2         { color: #e8e8e8; }
    .done-banner p          { color: #888; }
    .progress-track         { background: #333; }
    .progress-fill          { background: #ccc; }
}

/* ── Header ── */
h1 { font-size: 1.5rem; font-weight: 700; margin: 28px 0 4px; text-align: center; }
.subtitle { text-align: center; color: #888; margin-bottom: 20px; font-size: 0.9rem; }

/* ── Progress ── */
.progress-wrap  { margin-bottom: 16px; }
.progress-label { font-size: 0.8rem; color: #888; margin-bottom: 4px; text-align: right; }
.progress-track { height: 4px; border-radius: 4px; overflow: hidden; }
.progress-fill  { height: 100%; border-radius: 4px; transition: width 0.3s ease; }

/* ── Instructions ── */
.instructions {
    border: 1px solid;
    border-radius: 8px;
    padding: 16px 20px;
    margin-bottom: 16px;
    font-size: 0.88rem;
    line-height: 1.6;
}
.scale-list { list-style: none; padding: 0; margin: 10px 0 0; display: flex; flex-direction: column; gap: 6px; }
.scale-list li { display: flex; align-items: center; gap: 10px; font-size: 0.85rem; }
.scale-val {
    background: #222; color: #fff;
    border-radius: 4px; padding: 2px 8px;
    font-weight: 700; font-size: 0.8rem;
    min-width: 36px; text-align: center;
    flex-shrink: 0;
}

/* ── Done banner ── */
.done-banner {
    border: 1px solid;
    border-radius: 8px;
    padding: 64px 24px;
    text-align: center;
    margin: 32px 0;
}
.done-icon  { font-size: 3.5rem; margin-bottom: 16px; }
.done-banner h2 { font-size: 1.8rem; margin: 0 0 12px; }
.done-banner p  { margin: 0; font-size: 0.95rem; line-height: 1.7; }

footer { display: none !important; }
"""


# ── App ───────────────────────────────────────────────────────────────────────
def make_app():
    with gr.Blocks(css=CSS, title="Image Similarity Rating") as demo:

        # State
        user_id_state = gr.State(lambda: str(uuid.uuid4()))
        queue_state   = gr.State([])
        index_state   = gr.State(0)
        votes_state   = gr.State([])
        total_state   = gr.State(0)

        # Header
        gr.HTML("<h1>Image Similarity Rating</h1>")
        gr.HTML("<p class='subtitle'>Rate how similar Image B is to Image A.</p>")

        # Done banner — hidden until all pairs are rated
        done_html = gr.HTML(visible=False)

        # Rating UI — hidden when done
        with gr.Column(visible=True) as rating_col:

            # Progress
            progress_html = gr.HTML()

            # Images
            with gr.Row(equal_height=True):
                img_left  = gr.Image(label="Image A — Original",  show_label=True,
                                     interactive=False, height=520)
                img_right = gr.Image(label="Image B — Generated", show_label=True,
                                     interactive=False, height=520)

            # Instructions
            gr.HTML("""
            <div class="instructions">
                <strong>How similar is Image B to Image A?</strong><br>
                Image A is the original; Image B was reconstructed by an AI model.
                Rate their overall visual and semantic similarity:
                <ul class="scale-list">
                    <li><span class="scale-val">0</span>   Completely different</li>
                    <li><span class="scale-val">1–3</span> Very different; only a few elements match</li>
                    <li><span class="scale-val">4–6</span> Partial match; some key elements correct, but notable differences</li>
                    <li><span class="scale-val">7–9</span> Strong match; mostly correct with minor differences</li>
                    <li><span class="scale-val">10</span>  Identical or indistinguishable</li>
                </ul>
            </div>
            """)

            # Slider + button
            score_slider = gr.Slider(minimum=0, maximum=10, step=1, value=5,
                                     label="Similarity score (0–10)", interactive=True)
            next_btn = gr.Button("Submit and continue →", variant="primary", size="lg")

        # ── Helpers ───────────────────────────────────────────────────────

        def build_progress(idx, total):
            pct = int(idx / total * 100) if total else 0
            return f"""
            <div class="progress-wrap">
              <div class="progress-label">{idx} / {total}</div>
              <div class="progress-track">
                <div class="progress-fill" style="width:{pct}%"></div>
              </div>
            </div>"""

        # ── Init on load ──────────────────────────────────────────────────

        def on_load(user_id):
            queue = random.sample(_pairs, len(_pairs))
            total = len(queue)
            entry = queue[0]
            return (
                queue, 0, [], total,
                build_progress(0, total),
                entry["original_image_url"],
                entry["generated_image_url"],
                5,
                gr.update(visible=True),   # rating_col visible
                gr.update(visible=False),  # done_html hidden
            )

        demo.load(
            on_load,
            inputs=[user_id_state],
            outputs=[queue_state, index_state, votes_state, total_state,
                     progress_html, img_left, img_right, score_slider,
                     rating_col, done_html],
        )

        # ── Submit vote ───────────────────────────────────────────────────

        def on_next(score, queue, idx, votes, total, user_id):
            entry = queue[idx]
            vote = {
                "user_id":             user_id,
                "timestamp":           datetime.utcnow().isoformat(),
                "vote_index":          idx,
                "score":               int(score),
                "describer":           entry["describer"],
                "generator":           entry["generator"],
                "experiment":          entry["experiment"],
                "episode":             entry["episode"],
                "turn":                entry["turn"],
                "generated_image_url": entry["generated_image_url"],
                "original_image_url":  entry["original_image_url"],
            }
            votes = votes + [vote]
            idx  += 1

            # ── All pairs rated → show done banner, hide everything else ──
            if idx >= total:
                save_votes_to_hub(votes)
                done = """
                <div class="done-banner">
                  <div class="done-icon">✓</div>
                  <h2>Thank you!</h2>
                  <p>You have rated all image pairs.<br>
                     Your responses have been saved and will help us evaluate AI-generated images.</p>
                </div>"""
                return (
                    votes, idx,
                    build_progress(total, total),
                    gr.update(),   # img_left unchanged (hidden with column)
                    gr.update(),   # img_right unchanged
                    gr.update(),   # score_slider unchanged
                    gr.update(visible=False),  # rating_col → hide entire block
                    gr.update(value=done, visible=True),  # done_html → show
                )

            # ── Next pair ─────────────────────────────────────────────────
            next_entry = queue[idx]
            return (
                votes, idx,
                build_progress(idx, total),
                gr.update(value=next_entry["original_image_url"]),
                gr.update(value=next_entry["generated_image_url"]),
                5,
                gr.update(visible=True),   # rating_col stays visible
                gr.update(visible=False),  # done_html stays hidden
            )

        next_btn.click(
            on_next,
            inputs=[score_slider, queue_state, index_state, votes_state,
                    total_state, user_id_state],
            outputs=[votes_state, index_state,
                     progress_html, img_left, img_right, score_slider,
                     rating_col, done_html],
        )

    return demo


if __name__ == "__main__":
    make_app().launch()