#!/usr/bin/env python3
"""Streaming-ish web server for the from-scratch Zipformer-M CTC (Hindi/Hinglish).

Same browser/websocket protocol as the old NeMo server (mic -> 16k PCM int16 over /ws;
server returns {"type":"partial"/"final","text":...}; supports {cmd:reset|flush}).
Server accumulates audio and re-decodes with lhotse-Fbank + causal encoder + CTC greedy
(the same forward as eval_wer.py, so output matches reported WER). ▁ -> space.

Env: EXP_DIR, LANG_DIR, EPOCH, AVG (checkpoint averaging), PORT.
"""
import os, json, argparse, asyncio, math
import numpy as np
import torch
import sys
sys.path.insert(0, "/root/icefall/egs/hindi/ASR/zipformer"); sys.path.insert(0, "/root/icefall")
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from train import add_model_arguments, get_model, get_params
from icefall.lexicon import Lexicon
from icefall.checkpoint import average_checkpoints, load_checkpoint
from icefall.decode import ctc_greedy_search
from icefall.utils import make_pad_mask
from lhotse import Fbank, FbankConfig

EXP_DIR = os.environ.get("EXP_DIR", "/workspace/hindi_ft/asr_ctc/exp_p1")
LANG_DIR = os.environ.get("LANG_DIR", "/workspace/hindi_ft/asr_ctc/data/lang_char")
EPOCH = int(os.environ.get("EPOCH", "30"))
AVG = int(os.environ.get("AVG", "5"))
PORT = int(os.environ.get("PORT", "8080"))
SR = 16000
LOG_EPS = math.log(1e-10)
WB = "▁"
DEVICE = torch.device("cuda", 0) if torch.cuda.is_available() else torch.device("cpu")

print(f"[boot] loading zipformer CTC exp={EXP_DIR} epoch={EPOCH} avg={AVG}", flush=True)
_ap = argparse.ArgumentParser(); add_model_arguments(_ap)
params = get_params(); params.update(vars(_ap.parse_args([])))
params.causal = True; params.chunk_size = "16,32,64,-1"; params.left_context_frames = "64,128,256,-1"
params.use_ctc = True; params.use_transducer = False
lexicon = Lexicon(LANG_DIR)
params.blank_id = lexicon.token_table["<blk>"]; params.vocab_size = max(lexicon.tokens) + 1
model = get_model(params)
if AVG > 1:
    start = EPOCH - AVG + 1
    fns = [f"{EXP_DIR}/epoch-{e}.pt" for e in range(start, EPOCH + 1) if os.path.exists(f"{EXP_DIR}/epoch-{e}.pt")]
    print(f"[boot] averaging {len(fns)} ckpts", flush=True)
    model.load_state_dict(average_checkpoints(fns, device=DEVICE), strict=False)
else:
    load_checkpoint(f"{EXP_DIR}/epoch-{EPOCH}.pt", model)
model.to(DEVICE).eval()
fbank = Fbank(FbankConfig(num_mel_bins=80))
print(f"[boot] ready on {DEVICE}, vocab={params.vocab_size}", flush=True)


@torch.no_grad()
def transcribe(samples: np.ndarray) -> str:
    if samples.shape[0] < SR * 0.2:
        return ""
    feats = fbank.extract(torch.from_numpy(samples), SR)  # (T,80)
    feat = torch.as_tensor(np.asarray(feats), dtype=torch.float32).unsqueeze(0).to(DEVICE)
    flens = torch.tensor([feat.shape[1]], device=DEVICE) + 30
    feat = torch.nn.functional.pad(feat, (0, 0, 0, 30), value=LOG_EPS)
    x, xl = model.encoder_embed(feat, flens)
    mask = make_pad_mask(xl); x = x.permute(1, 0, 2)
    enc, el = model.encoder(x, xl, mask); enc = enc.permute(1, 0, 2)
    ctc = model.ctc_output(enc)
    ids = ctc_greedy_search(ctc, el)[0]
    return "".join(lexicon.token_table[i] for i in ids).replace(WB, " ").strip()


app = FastAPI()


class StreamState:
    def __init__(self):
        self.raw = np.zeros(0, dtype=np.float32)
        self.last = 0
        self.transcript = ""

    def add(self, pcm):
        self.raw = np.concatenate([self.raw, pcm])
        if self.raw.shape[0] > SR * 40:  # cap 40s
            self.raw = self.raw[-SR * 40:]

    def process(self):
        self.transcript = transcribe(self.raw)
        return self.transcript


@app.get("/health")
async def health():
    return {"status": "ok", "model": "zipformer-M-ctc-hindi", "epoch": EPOCH, "avg": AVG, "device": str(DEVICE)}


@app.websocket("/ws")
async def ws_endpoint(ws: WebSocket):
    await ws.accept()
    state = StreamState()
    step = SR // 2  # re-decode every ~0.5s of new audio
    try:
        while True:
            msg = await ws.receive()
            if "bytes" in msg and msg["bytes"] is not None:
                pcm = np.frombuffer(msg["bytes"], dtype=np.int16).astype(np.float32) / 32768.0
                state.add(pcm)
                if state.raw.shape[0] - state.last >= step:
                    state.last = state.raw.shape[0]
                    text = await asyncio.to_thread(state.process)
                    await ws.send_text(json.dumps({"type": "partial", "text": text}))
            elif "text" in msg and msg["text"] is not None:
                cmd = json.loads(msg["text"])
                if cmd.get("cmd") == "reset":
                    state = StreamState()
                    await ws.send_text(json.dumps({"type": "reset"}))
                elif cmd.get("cmd") == "flush":
                    text = await asyncio.to_thread(state.process)
                    await ws.send_text(json.dumps({"type": "final", "text": text}))
    except WebSocketDisconnect:
        pass
    except Exception as e:
        try:
            await ws.send_text(json.dumps({"type": "error", "text": str(e)}))
        except Exception:
            pass


@app.get("/")
async def index():
    return HTMLResponse(HTML_PAGE)


HTML_PAGE = r"""<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8"/>
<meta name="viewport" content="width=device-width, initial-scale=1"/>
<title>Hindi ASR — Zipformer CTC</title>
<style>
  :root { color-scheme: light dark; }
  * { box-sizing: border-box; }
  body { margin:0; font-family: -apple-system, system-ui, Segoe UI, Roboto, sans-serif;
         background:#0b0f14; color:#e6edf3; display:flex; min-height:100vh; }
  .wrap { margin:auto; width:min(820px, 92vw); padding:32px 0; }
  h1 { font-size:20px; font-weight:600; margin:0 0 4px; }
  .sub { color:#8b949e; font-size:13px; margin-bottom:24px; }
  .card { background:#121820; border:1px solid #222c38; border-radius:14px; padding:22px; }
  .controls { display:flex; gap:12px; align-items:center; margin-bottom:18px; }
  button { border:none; border-radius:10px; padding:11px 18px; font-size:14px; font-weight:600;
           cursor:pointer; transition:.15s; }
  #mic { background:#2ea043; color:#fff; }
  #mic.rec { background:#da3633; }
  #mic:disabled { opacity:.5; cursor:not-allowed; }
  #clear { background:#21262d; color:#e6edf3; }
  .dot { width:10px; height:10px; border-radius:50%; background:#3a3f45; display:inline-block; }
  .dot.on { background:#da3633; box-shadow:0 0 0 0 rgba(218,54,51,.7); animation:p 1.3s infinite; }
  @keyframes p { 0%{box-shadow:0 0 0 0 rgba(218,54,51,.6)} 70%{box-shadow:0 0 0 9px rgba(218,54,51,0)} }
  .status { font-size:12px; color:#8b949e; margin-left:auto; }
  #out { min-height:180px; font-size:20px; line-height:1.55; white-space:pre-wrap;
         padding:16px; background:#0b0f14; border-radius:10px; border:1px solid #222c38; }
  #out .cursor { color:#2ea043; animation:b 1s steps(1) infinite; }
  @keyframes b { 50%{opacity:0} }
  .foot { margin-top:14px; font-size:11px; color:#6e7681; }
  code { background:#21262d; padding:1px 6px; border-radius:5px; }
</style>
</head>
<body>
<div class="wrap">
  <h1>Streaming Speech-to-Text</h1>
  <div class="sub">Zipformer-M · char-CTC · from scratch (Hindi/Hinglish)</div>
  <div class="card">
    <div class="controls">
      <button id="mic">● Start talking</button>
      <button id="clear">Clear</button>
      <span class="dot" id="dot"></span>
      <span class="status" id="status">idle</span>
    </div>
    <div id="out"><span class="cursor">▍</span></div>
    <div class="foot">Mic runs at your device rate, downsampled to 16&nbsp;kHz and streamed as PCM.
      Speak naturally — text updates every chunk.</div>
  </div>
</div>
<script>
const micBtn = document.getElementById('mic');
const clearBtn = document.getElementById('clear');
const out = document.getElementById('out');
const dot = document.getElementById('dot');
const statusEl = document.getElementById('status');

let ws, audioCtx, source, processor, stream, recording = false;
let text = "";

function render() { out.innerHTML = (text ? escapeHtml(text) + " " : "") + '<span class="cursor">▍</span>'; }
function escapeHtml(s){ return s.replace(/[&<>]/g, c=>({'&':'&amp;','<':'&lt;','>':'&gt;'}[c])); }

function openWS() {
  return new Promise((resolve, reject) => {
    const proto = location.protocol === 'https:' ? 'wss' : 'ws';
    ws = new WebSocket(`${proto}://${location.host}/ws`);
    ws.binaryType = 'arraybuffer';
    ws.onopen = () => resolve();
    ws.onerror = (e) => reject(e);
    ws.onmessage = (ev) => {
      const m = JSON.parse(ev.data);
      if (m.type === 'partial' || m.type === 'final') { text = m.text || text; render(); }
      else if (m.type === 'reset') { text = ""; render(); }
      else if (m.type === 'error') { statusEl.textContent = 'error: ' + m.text; }
    };
  });
}

function downsample(buffer, inRate, outRate) {
  if (outRate === inRate) return buffer;
  const ratio = inRate / outRate;
  const outLen = Math.floor(buffer.length / ratio);
  const result = new Float32Array(outLen);
  let pos = 0;
  for (let i = 0; i < outLen; i++) {
    const start = Math.floor(i * ratio), end = Math.floor((i + 1) * ratio);
    let sum = 0, n = 0;
    for (let j = start; j < end && j < buffer.length; j++) { sum += buffer[j]; n++; }
    result[i] = n ? sum / n : buffer[start] || 0;
  }
  return result;
}

async function start() {
  statusEl.textContent = 'connecting…';
  await openWS();
  stream = await navigator.mediaDevices.getUserMedia({
    audio: { channelCount: 1, echoCancellation: true, noiseSuppression: true, autoGainControl: true }
  });
  audioCtx = new (window.AudioContext || window.webkitAudioContext)();
  source = audioCtx.createMediaStreamSource(stream);
  processor = audioCtx.createScriptProcessor(4096, 1, 1);
  source.connect(processor);
  processor.connect(audioCtx.destination);
  const inRate = audioCtx.sampleRate;
  processor.onaudioprocess = (e) => {
    if (!recording || ws.readyState !== WebSocket.OPEN) return;
    const ds = downsample(e.inputBuffer.getChannelData(0), inRate, 16000);
    const pcm = new Int16Array(ds.length);
    for (let i = 0; i < ds.length; i++) { let s = Math.max(-1, Math.min(1, ds[i])); pcm[i] = s * 32767; }
    ws.send(pcm.buffer);
  };
  recording = true;
  micBtn.classList.add('rec'); micBtn.textContent = '■ Stop';
  dot.classList.add('on'); statusEl.textContent = 'listening…';
}

function stop() {
  recording = false;
  if (processor) processor.disconnect();
  if (source) source.disconnect();
  if (stream) stream.getTracks().forEach(t => t.stop());
  if (audioCtx) audioCtx.close();
  if (ws && ws.readyState === WebSocket.OPEN) ws.send(JSON.stringify({cmd:'flush'}));
  micBtn.classList.remove('rec'); micBtn.textContent = '● Start talking';
  dot.classList.remove('on'); statusEl.textContent = 'stopped';
}

micBtn.onclick = async () => {
  micBtn.disabled = true;
  try { if (!recording) await start(); else stop(); }
  catch (e) { statusEl.textContent = 'mic error: ' + e.message; }
  micBtn.disabled = false;
};
clearBtn.onclick = () => {
  text = ""; render();
  if (ws && ws.readyState === WebSocket.OPEN) ws.send(JSON.stringify({cmd:'reset'}));
};
render();
</script>
</body>
</html>
"""


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=PORT)