import gradio as gr
import json
import os
import csv
import io
from datetime import datetime
from huggingface_hub import HfApi, hf_hub_download
import tempfile

HF_TOKEN = os.environ.get("HF_TOKEN", "")
DATASET_REPO = os.environ.get("DATASET_REPO", "")
RESULTS_FILE = "results.csv"

AUDIO_ITEMS = [
    {
        "id": "sample_01",
        "audio_path": "audios/sample1.wav",
        "transcript": "She opened the door and walked inside.",
        "stress_words": ["opened", "walked"],
    },
    {
        "id": "sample_02",
        "audio_path": "audios/sample2.wav",
        "transcript": "The storm destroyed every single house on the street.",
        "stress_words": ["destroyed", "every", "street"],
    },
    {
        "id": "sample_03",
        "audio_path": "audios/sample3.wav",
        "transcript": "I never said she stole the money.",
        "stress_words": ["never", "stole"],
    },
]

TOTAL = len(AUDIO_ITEMS)

_theme = gr.themes.Default(
    primary_hue="gray",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "sans-serif"],
)

_css = """
body, .gradio-container {
    font-family: 'Inter', sans-serif !important;
}
.gradio-container {
    max-width: 820px !important;
    margin: 0 auto;
}
#card, #done-card {
    background: #ffffff;
    border: 1px solid #e5e7eb;
    border-radius: 12px;
    padding: 24px;
}
#next-btn {
    background: #ffffff !important;
    color: #111827 !important;
    border: 1px solid #d1d5db !important;
    border-radius: 8px !important;
    font-weight: 600 !important;
}
#next-btn:hover {
    background: #f9fafb !important;
}
#result-box {
    border-radius: 10px;
}
"""


def highlight_transcript(transcript: str, stress_words: list) -> str:
    words = transcript.split()
    sw_lower = [w.lower().strip(".,!?;:") for w in stress_words]
    highlighted = []
    for w in words:
        clean = w.lower().strip(".,!?;:")
        if clean in sw_lower:
            highlighted.append(
                f'<span style="color:#dc2626;font-weight:700;">{w}</span>'
            )
        else:
            highlighted.append(w)
    return " ".join(highlighted)


def build_transcript_html(item: dict) -> str:
    stress_list = ", ".join(
        [f'<span style="color:#dc2626;font-weight:700;">{w}</span>' for w in item["stress_words"]]
    )
    body = highlight_transcript(item["transcript"], item["stress_words"])
    return (
        '<div style="background:#ffffff;border:1px solid #e5e7eb;border-radius:10px;'
        'padding:18px 20px;margin-bottom:8px;">'
        '<p style="margin:0 0 8px;font-size:0.82rem;color:#6b7280;'
        'text-transform:uppercase;letter-spacing:.05em;font-weight:600;">Transcript</p>'
        f'<p style="margin:0 0 14px;font-size:1.02rem;line-height:1.7;color:#111827;">{body}</p>'
        '<p style="margin:0;font-size:0.82rem;color:#6b7280;text-transform:uppercase;'
        'letter-spacing:.05em;font-weight:600;">Stress words</p>'
        f'<p style="margin:4px 0 0;font-size:0.95rem;color:#111827;">{stress_list}</p>'
        '</div>'
    )


def progress_html(idx: int) -> str:
    pct = int(idx / TOTAL * 100)
    return (
        '<div>'
        '<div style="display:flex;justify-content:space-between;margin-bottom:6px;">'
        f'<span style="font-size:.85rem;color:#6b7280;">Audio {idx} of {TOTAL}</span>'
        f'<span style="font-size:.85rem;color:#111827;font-weight:600;">{pct}%</span>'
        '</div>'
        '<div style="background:#e5e7eb;border-radius:999px;height:6px;">'
        f'<div style="background:#111827;width:{pct}%;height:6px;border-radius:999px;"></div>'
        '</div></div>'
    )


def push_results_to_hub(rows: list) -> str:
    if not HF_TOKEN or not DATASET_REPO:
        return "Results were not saved because HF_TOKEN or DATASET_REPO is missing."

    api = HfApi(token=HF_TOKEN)
    fieldnames = [
        "timestamp", "evaluator_id", "audio_id", "transcript",
        "stress_words", "naturalness", "stress_quality", "comment",
    ]

    existing_rows = []
    try:
        path = hf_hub_download(
            repo_id=DATASET_REPO,
            filename=RESULTS_FILE,
            repo_type="dataset",
            token=HF_TOKEN,
        )
        with open(path, newline="", encoding="utf-8") as f:
            existing_rows = list(csv.DictReader(f))
    except Exception:
        pass

    all_rows = existing_rows + rows
    buf = io.StringIO()
    writer = csv.DictWriter(buf, fieldnames=fieldnames)
    writer.writeheader()
    writer.writerows(all_rows)

    with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False, encoding="utf-8") as tmp:
        tmp.write(buf.getvalue())
        tmp_path = tmp.name

    api.upload_file(
        path_or_fileobj=tmp_path,
        path_in_repo=RESULTS_FILE,
        repo_id=DATASET_REPO,
        repo_type="dataset",
        commit_message=f"Add {len(rows)} evaluation(s) — {datetime.utcnow().strftime('%Y-%m-%d %H:%M UTC')}",
    )
    os.unlink(tmp_path)
    return f"Saved {len(rows)} responses to {DATASET_REPO}/{RESULTS_FILE}"


def make_app():
    with gr.Blocks(title="TTS Evaluation Survey") as demo:
        idx_state = gr.State(0)
        answers_state = gr.State([])

        gr.HTML("""
        <div style="text-align:center;padding:12px 0 8px;">
          <h1 style="margin:0;font-size:1.5rem;font-weight:700;color:#1d4ed8;">TTS Evaluation Survey</h1>
          <p style="margin:6px 0 0;color:#6b7280;font-size:.95rem;">
            Listen to each sample and rate naturalness and stress quality from 1 to 5.
          </p>
        </div>
        """)

        evaluator_id = gr.Textbox(
            label="Your name / evaluator ID",
            placeholder="e.g. listener_01",
            max_lines=1,
        )

        progress_bar = gr.HTML(value=progress_html(0))

        with gr.Group(elem_id="card", visible=True) as eval_card:
            audio_label = gr.HTML(
                value=(
                    f'<p style="font-size:.82rem;font-weight:600;color:#6b7280;'
                    f'text-transform:uppercase;letter-spacing:.05em;margin:0 0 8px;">'
                    f'Audio sample 1 / {TOTAL}</p>'
                )
            )
            transcript_box = gr.HTML(value=build_transcript_html(AUDIO_ITEMS[0]))
            audio_player = gr.Audio(
                value=AUDIO_ITEMS[0]["audio_path"],
                label="Listen",
                interactive=False,
            )
            naturalness = gr.Radio(
                choices=["1", "2", "3", "4", "5"],
                label="Naturalness (1 = poor, 5 = excellent)",
            )
            stress_quality = gr.Radio(
                choices=["1", "2", "3", "4", "5"],
                label="Stress quality (1 = poor, 5 = excellent)",
            )
            comment = gr.Textbox(
                label="Optional comment",
                placeholder="Any remarks about this sample",
                lines=2,
            )
            next_btn = gr.Button("Save and Next", elem_id="next-btn")

        with gr.Group(elem_id="done-card", visible=False) as done_card:
            done_html = gr.HTML(value="")

        def next_audio(idx, answers, eval_id, nat, stress, cmt):
            if nat is None or stress is None:
                gr.Warning("Please rate both Naturalness and Stress Quality.")
                return (
                    idx, answers, gr.update(), gr.update(), gr.update(), gr.update(),
                    nat, stress, cmt,
                    gr.update(visible=True),
                    gr.update(visible=False),
                    gr.update(value=""),
                )

            item = AUDIO_ITEMS[idx]
            record = {
                "timestamp": datetime.utcnow().isoformat(),
                "evaluator_id": eval_id.strip() if eval_id else "anonymous",
                "audio_id": item["id"],
                "transcript": item["transcript"],
                "stress_words": "|".join(item["stress_words"]),
                "naturalness": nat,
                "stress_quality": stress,
                "comment": cmt.strip() if cmt else "",
            }
            new_answers = answers + [record]
            new_idx = idx + 1

            if new_idx >= TOTAL:
                status = push_results_to_hub(new_answers)
                completion_html = (
    '<div style="text-align:center; background:#f8fafc; border:1px solid #dbeafe; '
    'border-radius:12px; padding:20px 16px;">'
    '<h2 style="margin:0 0 8px; color:#1e3a8a; font-size:1.2rem; font-weight:700;">'
    'Evaluation completed</h2>'
    '<p style="margin:0 0 10px; color:#334155; font-size:0.98rem;">'
    'All responses have been recorded. Please close the session.</p>'
    f'<p style="margin:0; color:#0f172a; font-size:.95rem; font-weight:500;"></p>'
    '</div>'
)
                print(json.dumps(new_answers, ensure_ascii=False, indent=2))
                return (
                    new_idx, new_answers, progress_html(new_idx),
                    gr.update(visible=False),
                    gr.update(visible=False),
                    gr.update(value=None),
                    None, None, "",
                    gr.update(visible=False),
                    gr.update(visible=True),
                    gr.update(value=completion_html),
                )

            next_item = AUDIO_ITEMS[new_idx]
            next_label = (
                f'<p style="font-size:.82rem;font-weight:600;color:#6b7280;'
                f'text-transform:uppercase;letter-spacing:.05em;margin:0 0 8px;">'
                f'Audio sample {new_idx + 1} / {TOTAL}</p>'
            )
            return (
                new_idx, new_answers, progress_html(new_idx), next_label,
                build_transcript_html(next_item), next_item["audio_path"],
                None, None, "",
                gr.update(visible=True),
                gr.update(visible=False),
                gr.update(value=""),
            )

        next_btn.click(
            fn=next_audio,
            inputs=[idx_state, answers_state, evaluator_id, naturalness, stress_quality, comment],
            outputs=[
                idx_state, answers_state, progress_bar,
                audio_label, transcript_box, audio_player,
                naturalness, stress_quality, comment,
                eval_card, done_card, done_html,
            ],
        )

    return demo


demo = make_app()

if __name__ == "__main__":
    demo.launch(theme=_theme, css=_css)