"""
maven dataset QA server
serves audio + text for review, writes keep/delete decisions back to txt file.

usage:
  pip install fastapi uvicorn
  python qa_server.py --wavs dataset/test --txt data/neutral-voice-dataset-list-clean.txt
  open http://localhost:8765
"""

import csv
import json
import argparse
import shutil
from pathlib import Path

from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
import uvicorn

# ── args ──────────────────────────────────────────────────────────────────────

parser = argparse.ArgumentParser()
parser.add_argument("--wavs", default="dataset/test", help="folder with wav files")
parser.add_argument("--txt",  required=True, help="txt file with sentences (one per line)")
parser.add_argument("--port", type=int, default=8765)
args = parser.parse_args()

WAV_DIR    = Path(args.wavs)
TXT_FILE   = Path(args.txt)
STATE_FILE = TXT_FILE.parent / (TXT_FILE.stem + "-qa-state.json")
META_FILE  = Path("tts/dataset/metadata.csv")

# ── state ─────────────────────────────────────────────────────────────────────

def load_sentences() -> list[str]:
    return [l.strip() for l in TXT_FILE.read_text(encoding="utf-8").splitlines() if l.strip()]

def load_state() -> dict:
    if STATE_FILE.exists():
        return json.loads(STATE_FILE.read_text())
    return {}  # index → "keep" | "delete"

def save_state(state: dict):
    STATE_FILE.write_text(json.dumps(state, indent=2, ensure_ascii=False))

def _build_text_to_wav() -> dict[str, str]:
    """Read metadata.csv (read-only) and build text → absolute wav path mapping."""
    mapping = {}
    if not META_FILE.exists():
        return mapping
    with open(META_FILE, "r", encoding="utf-8") as f:
        for row in csv.DictReader(f):
            text = row["text"].strip()
            file_path = row["file_path"].strip()
            # resolve to absolute path
            p = Path(file_path)
            if not p.is_absolute():
                # file_path is like "tts/dataset/neutral/0000.wav"
                p = Path(__file__).resolve().parent.parent / file_path
            if p.exists():
                mapping[text] = str(p)
    return mapping

# cache at startup
_TEXT_TO_WAV = _build_text_to_wav()

def wav_for(idx: int) -> Path | None:
    """Find wav file for sentence at txt line index via metadata.csv lookup."""
    sentences = load_sentences()
    if idx >= len(sentences):
        return None
    text = sentences[idx]
    wav_path = _TEXT_TO_WAV.get(text)
    if wav_path and Path(wav_path).exists():
        return Path(wav_path)
    return None

def get_deleted_indices() -> list[int]:
    """Get indices of deleted lines — from state first, then fallback to deleted/ folder."""
    state = load_state()
    indices = [int(k) for k, v in state.items() if v == "delete"]
    if indices:
        return sorted(indices)

    # fallback: scan deleted/ folder for numeric wav filenames
    deleted_dir = WAV_DIR / "deleted"
    if not deleted_dir.exists():
        return []
    for p in deleted_dir.glob("*.wav"):
        stem = p.stem.lstrip("0") or "0"
        if stem.isdigit():
            indices.append(int(stem))
    return sorted(indices)

# ── app ───────────────────────────────────────────────────────────────────────

app = FastAPI()

@app.get("/api/info")
def info():
    sentences = load_sentences()
    state     = load_state()
    wavs      = sorted(int(p.stem) for p in WAV_DIR.glob("*.wav") if p.stem.isdigit())

    reviewed  = len(state)
    kept      = sum(1 for v in state.values() if v == "keep")
    deleted   = sum(1 for v in state.values() if v == "delete")
    missing   = [i for i in range(len(sentences)) if wav_for(i) is None]

    return {
        "total":    len(sentences),
        "reviewed": reviewed,
        "kept":     kept,
        "deleted":  deleted,
        "missing":  missing,
        "wavs":     wavs,
    }

@app.get("/api/next")
def next_unreviewed(after: int = -1):
    sentences = load_sentences()
    state     = load_state()

    for i in range(after + 1, len(sentences)):
        if str(i) not in state and wav_for(i) is not None:
            return {
                "index": i,
                "text":  sentences[i],
                "total": len(sentences),
                "decision": state.get(str(i)),
            }

    return {"index": -1, "text": None, "total": len(sentences)}

@app.get("/api/sentence/{idx}")
def get_sentence(idx: int):
    sentences = load_sentences()
    state     = load_state()
    if idx >= len(sentences):
        raise HTTPException(404)
    return {
        "index":    idx,
        "text":     sentences[idx],
        "total":    len(sentences),
        "decision": state.get(str(idx)),
        "has_wav":  wav_for(idx) is not None,
    }

@app.get("/api/audio/{idx}")
def get_audio(idx: int):
    wav = wav_for(idx)
    if wav is None:
        raise HTTPException(404, "wav not found")
    return FileResponse(wav, media_type="audio/wav")

class Decision(BaseModel):
    decision: str  # "keep" | "delete"

@app.post("/api/decide/{idx}")
def decide(idx: int, body: Decision):
    if body.decision not in ("keep", "delete"):
        raise HTTPException(400, "decision must be keep or delete")

    sentences = load_sentences()
    if idx >= len(sentences):
        raise HTTPException(404)

    state = load_state()
    state[str(idx)] = body.decision
    save_state(state)
    return {"ok": True}

@app.post("/api/export")
def export():
    """write final txt with only kept sentences + list of missing wavs."""
    sentences = load_sentences()
    state     = load_state()

    kept    = []
    missing = []

    for i, s in enumerate(sentences):
        decision = state.get(str(i))
        if decision == "keep":
            kept.append(s)
        elif decision is None and wav_for(i) is None:
            missing.append(i)
        # deleted → skip

    out_path = TXT_FILE.parent / (TXT_FILE.stem + "-final.txt")
    out_path.write_text("\n".join(kept), encoding="utf-8")

    missing_path = TXT_FILE.parent / (TXT_FILE.stem + "-missing.json")
    missing_path.write_text(json.dumps(missing), encoding="utf-8")

    return {
        "kept":    len(kept),
        "missing": missing,
        "out":     str(out_path),
    }

# ── ui ────────────────────────────────────────────────────────────────────────

HTML = r"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Maven QA</title>
<style>
  @import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@300;400;500;700&display=swap');

  *, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }

  :root {
    --bg:      #0a0a0f;
    --surface: #111118;
    --border:  #1e1e2e;
    --text:    #c8c8d8;
    --dim:     #555570;
    --keep:    #3ddc84;
    --delete:  #ff4d6d;
    --accent:  #7c6af7;
    --font:    'JetBrains Mono', monospace;
  }

  body {
    background: var(--bg);
    color: var(--text);
    font-family: var(--font);
    min-height: 100vh;
    display: flex;
    flex-direction: column;
    align-items: center;
    justify-content: center;
    padding: 2rem;
  }

  header {
    position: fixed;
    top: 0; left: 0; right: 0;
    padding: 1rem 2rem;
    background: var(--bg);
    border-bottom: 1px solid var(--border);
    display: flex;
    align-items: center;
    justify-content: space-between;
    z-index: 10;
  }

  .logo { font-size: 0.75rem; color: var(--accent); letter-spacing: 0.2em; text-transform: uppercase; }

  .stats {
    display: flex;
    gap: 1.5rem;
    font-size: 0.7rem;
    color: var(--dim);
  }
  .stats span b { color: var(--text); }

  .card {
    width: 100%;
    max-width: 680px;
    margin-top: 4rem;
  }

  .index {
    font-size: 0.65rem;
    color: var(--dim);
    letter-spacing: 0.1em;
    margin-bottom: 0.75rem;
  }

  .sentence {
    font-size: 1.1rem;
    font-weight: 400;
    line-height: 1.6;
    color: var(--text);
    background: var(--surface);
    border: 1px solid var(--border);
    border-radius: 4px;
    padding: 1.25rem 1.5rem;
    margin-bottom: 1.5rem;
    min-height: 80px;
    display: flex;
    align-items: center;
  }

  .player {
    display: flex;
    align-items: center;
    gap: 1rem;
    margin-bottom: 1.5rem;
  }

  .play-btn {
    width: 44px; height: 44px;
    border-radius: 50%;
    border: 1px solid var(--border);
    background: var(--surface);
    color: var(--text);
    cursor: pointer;
    display: flex; align-items: center; justify-content: center;
    font-size: 1rem;
    flex-shrink: 0;
    transition: border-color 0.15s, background 0.15s;
  }
  .play-btn:hover { border-color: var(--accent); background: #1a1a2e; }
  .play-btn.playing { border-color: var(--accent); color: var(--accent); }

  .waveform {
    flex: 1;
    height: 44px;
    background: var(--surface);
    border: 1px solid var(--border);
    border-radius: 4px;
    display: flex;
    align-items: center;
    padding: 0 0.75rem;
    gap: 2px;
    overflow: hidden;
    cursor: pointer;
  }

  .bar {
    width: 2px;
    border-radius: 1px;
    background: var(--border);
    transition: background 0.1s;
    flex-shrink: 0;
  }
  .bar.active { background: var(--accent); }

  .actions {
    display: grid;
    grid-template-columns: 1fr 1fr;
    gap: 0.75rem;
    margin-bottom: 1rem;
  }

  .btn {
    padding: 0.875rem;
    border: 1px solid var(--border);
    background: var(--surface);
    color: var(--text);
    font-family: var(--font);
    font-size: 0.8rem;
    letter-spacing: 0.08em;
    text-transform: uppercase;
    cursor: pointer;
    border-radius: 4px;
    transition: all 0.15s;
    display: flex; align-items: center; justify-content: center; gap: 0.5rem;
  }

  .btn-keep  { border-color: var(--keep);   color: var(--keep); }
  .btn-keep:hover  { background: rgba(61,220,132,0.08); }
  .btn-keep.active  { background: rgba(61,220,132,0.15); }

  .btn-delete { border-color: var(--delete); color: var(--delete); }
  .btn-delete:hover { background: rgba(255,77,109,0.08); }
  .btn-delete.active { background: rgba(255,77,109,0.15); }

  .btn-nav {
    width: 100%;
    margin-top: 0.5rem;
    color: var(--dim);
    font-size: 0.72rem;
  }
  .btn-nav:hover { color: var(--text); border-color: var(--dim); }

  .btn-export {
    width: 100%;
    margin-top: 2rem;
    color: var(--accent);
    border-color: var(--accent);
    font-size: 0.72rem;
  }
  .btn-export:hover { background: rgba(124,106,247,0.08); }

  .progress-bar {
    height: 2px;
    background: var(--border);
    border-radius: 1px;
    margin-bottom: 2rem;
    overflow: hidden;
  }
  .progress-fill {
    height: 100%;
    background: var(--accent);
    border-radius: 1px;
    transition: width 0.3s ease;
  }

  .decision-badge {
    display: inline-block;
    font-size: 0.65rem;
    letter-spacing: 0.12em;
    text-transform: uppercase;
    padding: 0.2rem 0.5rem;
    border-radius: 2px;
    margin-bottom: 0.75rem;
  }
  .decision-badge.keep   { background: rgba(61,220,132,0.15); color: var(--keep); }
  .decision-badge.delete { background: rgba(255,77,109,0.15); color: var(--delete); }
  .decision-badge.none   { display: none; }

  .no-wav {
    font-size: 0.7rem;
    color: var(--delete);
    letter-spacing: 0.1em;
    margin-bottom: 0.75rem;
  }

  .export-result {
    margin-top: 1rem;
    font-size: 0.72rem;
    color: var(--keep);
    text-align: center;
  }

  kbd {
    display: inline-block;
    padding: 0.1rem 0.35rem;
    border: 1px solid var(--border);
    border-radius: 2px;
    font-size: 0.65rem;
    color: var(--dim);
    margin-left: 0.25rem;
  }
</style>
</head>
<body>

<header>
  <div class="logo">Maven // Dataset QA</div>
  <div class="stats">
    <span><b id="s-reviewed">0</b> reviewed</span>
    <span><b id="s-kept" style="color:var(--keep)">0</b> kept</span>
    <span><b id="s-deleted" style="color:var(--delete)">0</b> deleted</span>
    <span><b id="s-total">0</b> total</span>
  </div>
</header>

<div class="card">
  <div class="progress-bar"><div class="progress-fill" id="progress" style="width:0%"></div></div>

  <div class="index" id="idx">— / —</div>
  <div class="decision-badge none" id="badge"></div>
  <div class="no-wav" id="no-wav" style="display:none">⚠ no audio file</div>

  <div class="sentence" id="sentence">loading...</div>

  <div class="player">
    <button class="play-btn" id="play-btn" onclick="togglePlay()">▶</button>
    <div class="waveform" id="waveform" onclick="seekClick(event)"></div>
  </div>

  <div class="actions">
    <button class="btn btn-keep"   id="btn-keep"   onclick="decide('keep')">✓ keep<kbd>k</kbd></button>
    <button class="btn btn-delete" id="btn-delete"  onclick="decide('delete')">✕ delete<kbd>d</kbd></button>
  </div>

  <button class="btn btn-nav" onclick="navigate(-1)">← prev<kbd>←</kbd></button>
  <button class="btn btn-nav" style="margin-top:0.5rem" onclick="navigate(1)">next →<kbd>→</kbd></button>

  <button class="btn btn-export" onclick="doExport()">⬇ export final dataset</button>
  <button class="btn btn-export" style="margin-top:0.5rem;color:var(--delete);border-color:var(--delete)" onclick="doCleanup()">⊗ move deleted wavs to deleted/</button>
  <button class="btn btn-export" style="margin-top:0.5rem;color:var(--dim);border-color:var(--dim)" onclick="doResetDeleted()">↺ reset deleted → unreviewed (if wav exists)</button>
  <div class="export-result" id="export-result"></div>
</div>

<audio id="audio"></audio>

<script>
let currentIdx = 0;
let totalSentences = 0;
let audio = document.getElementById('audio');
let bars = [];
let animFrame = null;

// ── waveform ──────────────────────────────────────────────────────────────────

function buildWaveform() {
  const wf = document.getElementById('waveform');
  wf.innerHTML = '';
  bars = [];
  const n = Math.floor(wf.offsetWidth / 4);
  for (let i = 0; i < n; i++) {
    const b = document.createElement('div');
    b.className = 'bar';
    const h = 6 + Math.random() * 28;
    b.style.height = h + 'px';
    wf.appendChild(b);
    bars.push(b);
  }
}

function updateWaveform() {
  if (!bars.length) return;
  const pct = audio.duration ? audio.currentTime / audio.duration : 0;
  const active = Math.floor(pct * bars.length);
  bars.forEach((b, i) => b.classList.toggle('active', i < active));
  if (!audio.paused) animFrame = requestAnimationFrame(updateWaveform);
}

function seekClick(e) {
  if (!audio.src || !audio.duration) return;
  const rect = e.currentTarget.getBoundingClientRect();
  const pct  = (e.clientX - rect.left) / rect.width;
  audio.currentTime = pct * audio.duration;
  updateWaveform();
}

// ── playback ──────────────────────────────────────────────────────────────────

function togglePlay() {
  if (audio.paused) { audio.play(); }
  else              { audio.pause(); }
}

audio.addEventListener('play',  () => {
  document.getElementById('play-btn').classList.add('playing');
  document.getElementById('play-btn').textContent = '⏸';
  animFrame = requestAnimationFrame(updateWaveform);
});
audio.addEventListener('pause', () => {
  document.getElementById('play-btn').classList.remove('playing');
  document.getElementById('play-btn').textContent = '▶';
  cancelAnimationFrame(animFrame);
});
audio.addEventListener('ended', () => {
  document.getElementById('play-btn').textContent = '▶';
  bars.forEach(b => b.classList.remove('active'));
});

// ── load sentence ─────────────────────────────────────────────────────────────

async function loadSentence(idx) {
  const res  = await fetch(`/api/sentence/${idx}`);
  if (!res.ok) return;
  const data = await res.json();

  currentIdx      = idx;
  totalSentences  = data.total;

  document.getElementById('idx').textContent      = `${idx + 1} / ${data.total}`;
  document.getElementById('sentence').textContent = data.text || '—';

  // decision badge
  const badge = document.getElementById('badge');
  badge.className = 'decision-badge ' + (data.decision || 'none');
  badge.textContent = data.decision || '';

  // button active states
  document.getElementById('btn-keep').classList.toggle('active', data.decision === 'keep');
  document.getElementById('btn-delete').classList.toggle('active', data.decision === 'delete');

  // wav
  const noWav = document.getElementById('no-wav');
  if (data.has_wav) {
    noWav.style.display = 'none';
    audio.src = `/api/audio/${idx}?t=${Date.now()}`;
    audio.load();
    buildWaveform();
    // autoplay
    try { await audio.play(); } catch {}
  } else {
    noWav.style.display = 'block';
    audio.src = '';
  }

  await refreshStats();
}

// ── stats ─────────────────────────────────────────────────────────────────────

async function refreshStats() {
  const res  = await fetch('/api/info');
  const data = await res.json();
  document.getElementById('s-reviewed').textContent = data.reviewed;
  document.getElementById('s-kept').textContent     = data.kept;
  document.getElementById('s-deleted').textContent  = data.deleted;
  document.getElementById('s-total').textContent    = data.total;
  const pct = data.total ? (data.reviewed / data.total * 100) : 0;
  document.getElementById('progress').style.width   = pct + '%';
}

// ── decisions ─────────────────────────────────────────────────────────────────

async function decide(decision) {
  await fetch(`/api/decide/${currentIdx}`, {
    method: 'POST',
    headers: {'Content-Type': 'application/json'},
    body: JSON.stringify({decision}),
  });
  // auto-advance on decide
  navigate(1);
}

// ── navigation ────────────────────────────────────────────────────────────────

async function navigate(dir) {
  audio.pause();
  if (dir === 1) {
    // find next with wav
    const res  = await fetch(`/api/next?after=${currentIdx}`);
    const data = await res.json();
    if (data.index >= 0) loadSentence(data.index);
  } else {
    if (currentIdx > 0) loadSentence(currentIdx - 1);
  }
}

// ── export ────────────────────────────────────────────────────────────────────

async function doExport() {
  const res  = await fetch('/api/export', {method: 'POST'});
  const data = await res.json();
  const el   = document.getElementById('export-result');
  el.textContent = `✓ exported ${data.kept} sentences. missing wavs: ${data.missing.length}`;
}

// ── cleanup ───────────────────────────────────────────────────────────────────

async function doCleanup() {
  if (!confirm('Move all deleted wavs to deleted/ subfolder?')) return;
  const res  = await fetch('/api/cleanup', {method: 'POST'});
  const data = await res.json();
  const el   = document.getElementById('export-result');
  el.textContent = `⊗ moved ${data.moved} wavs to deleted/. errors: ${data.errors.length}`;
  el.style.color = 'var(--delete)';
}

async function doResetDeleted() {
  if (!confirm('Reset all deleted entries? This will clear them from the qa-state.json so they become unreviewed again.')) return;
  const res  = await fetch('/api/reset_deleted', {method: 'POST'});
  const data = await res.json();
  const el   = document.getElementById('export-result');
  el.style.color = 'var(--text)';
  el.textContent = `↺ reset ${data.cleared} entries — ${data.lines_removed} lines removed from txt, state cleaned`;
  await refreshStats();
  loadSentence(currentIdx);
}

// ── keyboard ──────────────────────────────────────────────────────────────────

document.addEventListener('keydown', e => {
  if (e.key === 'k')           decide('keep');
  else if (e.key === 'd')      decide('delete');
  else if (e.key === 'ArrowRight') navigate(1);
  else if (e.key === 'ArrowLeft')  navigate(-1);
  else if (e.key === ' ')      { e.preventDefault(); togglePlay(); }
});

// ── init ──────────────────────────────────────────────────────────────────────

(async () => {
  const res  = await fetch('/api/next?after=-1');
  const data = await res.json();
  const start = data.index >= 0 ? data.index : 0;
  loadSentence(start);
})();
</script>
</body>
</html>"""

@app.post("/api/reset_deleted")
def reset_deleted():
    """clear all delete decisions from the qa-state.json and remove corresponding lines from the txt file."""
    deleted_indices = get_deleted_indices()

    if deleted_indices:
        # Remove lines from txt file by index (descending to preserve indices)
        sentences = load_sentences()
        for idx in sorted(deleted_indices, reverse=True):
            if idx < len(sentences):
                del sentences[idx]

        TXT_FILE.write_text("\n".join(sentences) + "\n" if sentences else "", encoding="utf-8")

    # Clear all delete entries from state
    state = load_state()
    cleared = []
    for idx_str, decision in list(state.items()):
        if decision == "delete":
            del state[idx_str]
            cleared.append(int(idx_str))
    save_state(state)

    return {"cleared": len(cleared), "lines_removed": len(deleted_indices), "indices": deleted_indices}

@app.post("/api/cleanup")
def cleanup():
    """move deleted wavs to deleted/ subfolder."""
    state   = load_state()
    deleted_dir = WAV_DIR / "deleted"
    deleted_dir.mkdir(exist_ok=True)

    moved  = []
    errors = []

    for idx_str, decision in state.items():
        if decision != "delete":
            continue
        wav = wav_for(int(idx_str))
        if wav is None:
            continue
        dest = deleted_dir / wav.name
        try:
            shutil.move(str(wav), str(dest))
            moved.append(int(idx_str))
        except Exception as e:
            errors.append({"idx": int(idx_str), "error": str(e)})

    return {"moved": len(moved), "errors": errors, "dest": str(deleted_dir)}

@app.get("/", response_class=HTMLResponse)
def ui():
    return HTML

# ── run ───────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    print(f"[*] wavs:  {WAV_DIR}")
    print(f"[*] txt:   {TXT_FILE}")
    print(f"[*] state: {STATE_FILE}")
    print(f"[*] open   http://localhost:{args.port}")
    uvicorn.run(app, host="0.0.0.0", port=args.port, log_level="debug")