#!/usr/bin/env python3
"""TextStream — Live speech-to-text streaming on Apple Silicon.

Streams live microphone audio through Qwen3-ASR, displays results in a
browser (SSE), optionally pushes Grafana annotations, and saves daily
transcripts to ~/Documents/textstream/transcripts/YYYY-MM-DD/.

Usage:
    textstream                          # Start with Qwen3-ASR 0.6B (default)
    textstream --engine qwen-1.7b       # Use Qwen3-ASR 1.7B (higher accuracy)
    textstream --vad-threshold 0.5      # Stricter voice activity detection
    textstream --interval 2.5           # Seconds between streaming updates
    textstream --port 8080              # Custom port
"""

import sys
import os
import json
import time
import signal
import threading
import queue
import webbrowser
import argparse
import urllib.request
import urllib.error
from abc import ABC, abstractmethod
from pathlib import Path
from socketserver import ThreadingMixIn
from http.server import HTTPServer, BaseHTTPRequestHandler

import numpy as np
import sounddevice as sd

# ── Config ────────────────────────────────────────────────────────────────────
SAMPLE_RATE = 16000
DEFAULT_PORT = 7890
DEFAULT_INTERVAL = 2.5
DEFAULT_VAD_THRESHOLD = 0.4

GRAFANA_URL = os.environ.get("GRAFANA_URL", "https://triscient.grafana.net")
GRAFANA_TOKEN = os.environ.get("GRAFANA_SERVICE_ACCOUNT_TOKEN", "")

# ── State ─────────────────────────────────────────────────────────────────────
audio_chunks = []
buffer_lock = threading.Lock()
subscribers = []
sub_lock = threading.Lock()
running = True
paused = False
current_engine = None
engine_lock = threading.Lock()


def log(msg):
    print(f"\033[90m[textstream]\033[0m {msg}", file=sys.stderr, flush=True)


# ── Audio Capture ─────────────────────────────────────────────────────────────
def audio_callback(indata, frames, time_info, status):
    if status:
        log(f"audio warning: {status}")
    with buffer_lock:
        audio_chunks.append(indata[:, 0].copy())


def drain_buffer():
    global audio_chunks
    with buffer_lock:
        if not audio_chunks:
            return None
        chunk = np.concatenate(audio_chunks)
        audio_chunks = []
    return chunk


# ── SSE Broadcast ─────────────────────────────────────────────────────────────
def broadcast(event_data):
    payload = f"data: {json.dumps(event_data)}\n\n".encode()
    with sub_lock:
        dead = []
        for q in subscribers:
            try:
                q.put_nowait(payload)
            except Exception:
                dead.append(q)
        for d in dead:
            subscribers.remove(d)


# ── Grafana Annotation Push ───────────────────────────────────────────────────
def push_annotation(text):
    try:
        payload = json.dumps({
            "text": text,
            "tags": ["textstream"],
            "time": int(time.time() * 1000),
        }).encode()
        req = urllib.request.Request(
            f"{GRAFANA_URL}/api/annotations",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {GRAFANA_TOKEN}",
                "User-Agent": "textstream/1.0",
            },
        )
        urllib.request.urlopen(req, timeout=5)
    except Exception as e:
        log(f"grafana push failed: {e}")


# ── Transcript File Persistence ───────────────────────────────────────────────
TRANSCRIPT_DIR = Path.home() / "Documents" / "textstream" / "transcripts"
_transcript_file = None
_transcript_lock = threading.Lock()


def get_transcript_file():
    global _transcript_file
    from datetime import datetime

    now = datetime.now()
    day_dir = TRANSCRIPT_DIR / now.strftime("%Y-%m-%d")
    day_dir.mkdir(parents=True, exist_ok=True)

    if _transcript_file is None or _transcript_file.closed:
        session_name = now.strftime("%H-%M-%S")
        path = day_dir / f"{session_name}.txt"
        _transcript_file = open(path, "a", encoding="utf-8")
        log(f"Transcript: {path}")
    return _transcript_file


def save_transcript(text):
    from datetime import datetime

    with _transcript_lock:
        f = get_transcript_file()
        ts = datetime.now().strftime("%H:%M:%S")
        f.write(f"[{ts}] {text}\n")
        f.flush()


# ── Engine Abstraction ────────────────────────────────────────────────────────
class ASREngine(ABC):
    name: str = "unknown"
    description: str = ""

    @abstractmethod
    def start(self):
        """Initialize streaming state."""

    @abstractmethod
    def feed(self, audio_np: np.ndarray) -> tuple[str, str]:
        """Feed audio chunk, return (stable_text, draft_text)."""

    @abstractmethod
    def stop(self):
        """Finalize and clean up streaming state."""

    @abstractmethod
    def needs_manual_reset(self) -> bool:
        """Whether this engine needs periodic drift resets."""

    def reset(self):
        """Reset streaming state (for engines that drift)."""
        self.stop()
        self.start()


class QwenEngine(ASREngine):
    description = "Qwen3-ASR 0.6B — accurate (~2.3% WER), built-in context window"

    # Known hallucination patterns (model regurgitates its chat system prompt on noise)
    HALLUCINATION_PATTERNS = [
        "you are a helpful assistant",
        "i am a helpful assistant",
        "as an ai",
        "as a language model",
    ]

    def __init__(self, model_size="0.6b"):
        self.model_size = model_size
        self.name = "qwen" if model_size == "0.6b" else "qwen-1.7b"
        self.session = None
        self._state = None
        self._prev_stable_len = 0
        self._hallucination_streak = 0

    def load(self):
        model_map = {
            "0.6b": "Qwen/Qwen3-ASR-0.6B",
            "1.7b": "Qwen/Qwen3-ASR-1.7B",
        }
        log(f"Loading Qwen3-ASR {self.model_size}...")
        from mlx_qwen3_asr import Session
        self.session = Session(model=model_map[self.model_size])
        log("Qwen3-ASR ready")

    def start(self):
        if self.session is None:
            self.load()
        self._state = self.session.init_streaming(
            chunk_size_sec=2.0,
            max_context_sec=30.0,
            finalization_mode="accuracy",
            sample_rate=SAMPLE_RATE,
            unfixed_chunk_num=2,
            unfixed_token_num=5,
            language="English",
        )
        self._prev_stable_len = 0

    def _is_hallucination(self, text):
        lower = text.lower().strip()
        return any(p in lower for p in self.HALLUCINATION_PATTERNS)

    def feed(self, audio_np):
        self._state = self.session.feed_audio(audio_np, self._state)

        stable = self._state.stable_text.strip()
        full = self._state.text.strip()

        # Draft = everything after stable
        draft = ""
        if len(full) > len(stable):
            draft = full[len(stable):].strip()

        # Detect hallucination — model regurgitates system prompt on noise/music
        if self._is_hallucination(draft) or self._is_hallucination(full[-60:] if full else ""):
            self._hallucination_streak += 1
            if self._hallucination_streak >= 2:
                log("  -- hallucination detected, resetting stream --")
                self.stop()
                self.start()
                return "", ""
            return stable, ""  # suppress the hallucinated draft
        else:
            self._hallucination_streak = 0

        # Track newly finalized text for persistence
        new_text = ""
        if len(stable) > self._prev_stable_len:
            new_text = stable[self._prev_stable_len:].strip()
            # Don't persist hallucinated text
            if new_text and not self._is_hallucination(new_text):
                save_transcript(new_text)
                threading.Thread(
                    target=push_annotation, args=(new_text,), daemon=True
                ).start()
            self._prev_stable_len = len(stable)

        return stable, draft

    def stop(self):
        if self._state and self.session:
            try:
                self._state = self.session.finish_streaming(self._state)
                final = self._state.stable_text.strip()
                if len(final) > self._prev_stable_len:
                    remaining = final[self._prev_stable_len:].strip()
                    if remaining:
                        save_transcript(remaining)
                        threading.Thread(
                            target=push_annotation, args=(remaining,), daemon=True
                        ).start()
            except Exception as e:
                log(f"qwen finish error: {e}")
        self._state = None
        self._prev_stable_len = 0

    def needs_manual_reset(self):
        return False  # built-in 30s sliding window handles context


ENGINES = {
    "qwen": lambda: QwenEngine("0.6b"),
    "qwen-1.7b": lambda: QwenEngine("1.7b"),
}


# ── Streaming Transcription Loop ─────────────────────────────────────────────
SILENCE_STREAK_RESET = 4


def transcription_loop(interval, vad_threshold):
    global current_engine

    from .vad import contains_speech

    broadcast({"type": "status", "content": "Listening..."})
    log("Listening - speak into your microphone (streaming mode)")

    with engine_lock:
        engine = current_engine

    engine.start()

    update_count = 0
    total_audio_sec = 0.0
    total_infer_ms = 0.0
    prev_text = ""
    silence_streak = 0

    while running:
        time.sleep(interval)
        if paused:
            drain_buffer()
            continue

        # Check for engine switch
        with engine_lock:
            if current_engine is not engine:
                log("Engine switch detected, transitioning...")
                engine.stop()
                engine = current_engine
                engine.start()
                prev_text = ""
                silence_streak = 0
                update_count = 0
                total_audio_sec = 0.0
                total_infer_ms = 0.0
                broadcast({"type": "status", "content": f"Switched to {engine.name}"})
                continue

        chunk = drain_buffer()
        if chunk is None or len(chunk) < 800:
            silence_streak += 1
            if silence_streak >= SILENCE_STREAK_RESET and engine.needs_manual_reset():
                engine.reset()
                silence_streak = 0
            continue

        if not contains_speech(chunk, threshold=vad_threshold):
            silence_streak += 1
            if silence_streak >= SILENCE_STREAK_RESET and engine.needs_manual_reset():
                engine.reset()
                silence_streak = 0
            continue

        silence_streak = 0
        audio_sec = len(chunk) / SAMPLE_RATE

        t0 = time.perf_counter()
        try:
            stable, draft = engine.feed(chunk)
        except Exception as e:
            log(f"engine error: {e}")
            try:
                engine.reset()
            except Exception:
                pass
            continue
        elapsed_ms = (time.perf_counter() - t0) * 1000

        update_count += 1
        total_audio_sec += audio_sec
        total_infer_ms += elapsed_ms

        full_text = f"{stable} {draft}".strip()
        if not full_text or full_text == prev_text:
            continue
        prev_text = full_text

        broadcast({
            "type": "stream",
            "finalized": stable,
            "draft": draft,
        })

        avg_ms = total_infer_ms / update_count
        rtf = (total_infer_ms / 1000) / total_audio_sec if total_audio_sec > 0 else 0
        log(
            f"  [{elapsed_ms:.0f}ms | avg {avg_ms:.0f}ms | "
            f"RTF={rtf:.4f} | {engine.name}] "
            f"{stable[-50:]}|{draft[:35]}"
        )

    # Clean shutdown
    engine.stop()


# ── HTML ──────────────────────────────────────────────────────────────────────
HTML = r"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>TextStream</title>
<style>
*{margin:0;padding:0;box-sizing:border-box}
html,body{height:100%;overflow:hidden}
body{
  background:#09090b;
  color:#e8e8e8;
  font-family:-apple-system,BlinkMacSystemFont,'SF Pro Text','Inter',sans-serif;
  display:flex;flex-direction:column;
}
.topbar{
  display:flex;align-items:center;justify-content:space-between;
  padding:0.8rem 1.5rem;
  background:#111;
  border-bottom:1px solid #1a1a1a;
  flex-shrink:0;
  user-select:none;
}
.topbar .left{display:flex;align-items:center;gap:0.6rem}
.topbar .title{font-size:0.7rem;color:#555;letter-spacing:0.12em;text-transform:uppercase;font-weight:600}
.indicator{width:6px;height:6px;border-radius:50%;background:#34d399;transition:background 0.3s}
.indicator.off{background:#555;animation:none}
.indicator.on{animation:pulse 2.5s ease-in-out infinite}
@keyframes pulse{0%,100%{opacity:0.3}50%{opacity:1}}
.topbar .controls{display:flex;gap:0.5rem;align-items:center}
.btn{
  background:#1a1a1a;border:1px solid #2a2a2a;color:#888;
  padding:0.35rem 0.9rem;border-radius:4px;font-size:0.65rem;
  cursor:pointer;letter-spacing:0.06em;text-transform:uppercase;
  font-weight:500;transition:all 0.2s;
}
.btn:hover{background:#222;color:#ccc;border-color:#444}
.btn.active{background:#1a2e1a;border-color:#2a4a2a;color:#4ade80}
.btn.danger:hover{background:#2e1a1a;border-color:#4a2a2a;color:#f87171}
select.engine-select{
  background:#1a1a1a;border:1px solid #2a2a2a;color:#888;
  padding:0.35rem 0.5rem;border-radius:4px;font-size:0.65rem;
  cursor:pointer;letter-spacing:0.04em;
  font-weight:500;outline:none;
  -webkit-appearance:none;appearance:none;
}
select.engine-select:hover{background:#222;color:#ccc;border-color:#444}
select.engine-select:focus{border-color:#4ade80}
.content{
  flex:1;display:flex;align-items:flex-end;
  overflow:hidden;
}
.wrap{
  width:100%;max-width:900px;
  padding:0 2.5rem 3rem;
  margin:0 auto;
}
#text{
  font-size:1.65rem;
  font-weight:350;
  line-height:1.7;
  letter-spacing:-0.01em;
  text-align:left;
  word-wrap:break-word;
}
#text .draft{color:#555}
#status{font-size:1rem;color:#444;font-weight:300;padding-bottom:0.8rem}
</style>
</head>
<body>
<div class="topbar">
  <div class="left">
    <div class="indicator on" id="ind"></div>
    <span class="title">textstream</span>
  </div>
  <div class="controls">
    <select class="engine-select" id="engineSelect" onchange="switchEngine(this.value)">
      <option value="qwen">Qwen3 0.6B</option>
      <option value="qwen-1.7b">Qwen3 1.7B</option>
    </select>
    <button class="btn active" id="toggleBtn" onclick="togglePause()">Listening</button>
    <button class="btn danger" onclick="stopServer()">Stop</button>
  </div>
</div>
<div class="content">
  <div class="wrap">
    <div id="status">Loading model...</div>
    <div id="text"></div>
  </div>
</div>
<script>
const textEl=document.getElementById('text');
const statusEl=document.getElementById('status');
const toggleBtn=document.getElementById('toggleBtn');
const ind=document.getElementById('ind');
const engineSelect=document.getElementById('engineSelect');
let isPaused=false;
const MAX_CHARS=500;

function togglePause(){
  isPaused=!isPaused;
  fetch(isPaused?'/pause':'/resume');
  toggleBtn.textContent=isPaused?'Paused':'Listening';
  toggleBtn.classList.toggle('active',!isPaused);
  ind.classList.toggle('on',!isPaused);
  ind.classList.toggle('off',isPaused);
}
function stopServer(){
  if(confirm('Stop TextStream?')){
    fetch('/stop');
    document.body.innerHTML='<div style="display:flex;align-items:center;justify-content:center;height:100vh;color:#444;font-size:1.2rem">TextStream stopped.</div>';
  }
}
function switchEngine(eng){
  statusEl.style.display='block';
  statusEl.textContent='Switching to '+eng+'...';
  textEl.innerHTML='';
  fetch('/switch?engine='+encodeURIComponent(eng));
}

const src=new EventSource('/stream');
src.onmessage=e=>{
  const d=JSON.parse(e.data);
  if(d.type==='stream'){
    statusEl.style.display='none';
    const fin=d.finalized||'';
    const draft=d.draft||'';
    let visFin=fin;
    if(visFin.length>MAX_CHARS){
      const cut=visFin.indexOf(' ',visFin.length-MAX_CHARS);
      visFin=cut>0?visFin.slice(cut+1):visFin.slice(-MAX_CHARS);
    }
    if(draft){
      textEl.innerHTML=esc(visFin)+'<span class="draft"> '+esc(draft)+'</span>';
    }else{
      textEl.textContent=visFin;
    }
  }else if(d.type==='status'){
    statusEl.style.display='block';
    statusEl.textContent=d.content;
  }else if(d.type==='engine'){
    engineSelect.value=d.engine;
  }
};
src.onerror=()=>{
  statusEl.style.display='block';
  statusEl.textContent='Reconnecting...';
};
function esc(s){
  const d=document.createElement('div');
  d.textContent=s;
  return d.innerHTML;
}

// Set initial engine from server
fetch('/engine').then(r=>r.json()).then(d=>{engineSelect.value=d.engine});
</script>
</body>
</html>"""


# ── HTTP Server ───────────────────────────────────────────────────────────────
class ThreadedServer(ThreadingMixIn, HTTPServer):
    daemon_threads = True
    allow_reuse_address = True


class Handler(BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path == "/":
            self.send_response(200)
            self.send_header("Content-Type", "text/html; charset=utf-8")
            self.end_headers()
            self.wfile.write(HTML.encode())

        elif self.path == "/pause":
            global paused
            paused = True
            self.send_response(200)
            self.send_header("Content-Type", "text/plain")
            self.end_headers()
            self.wfile.write(b"paused")
            log("Paused")

        elif self.path == "/resume":
            paused = False
            self.send_response(200)
            self.send_header("Content-Type", "text/plain")
            self.end_headers()
            self.wfile.write(b"resumed")
            log("Resumed")

        elif self.path == "/stop":
            self.send_response(200)
            self.send_header("Content-Type", "text/plain")
            self.end_headers()
            self.wfile.write(b"TextStream stopped.\n")
            log("Stop requested via /stop endpoint")
            threading.Thread(target=lambda: os.kill(os.getpid(), signal.SIGTERM), daemon=True).start()

        elif self.path.startswith("/switch"):
            global current_engine
            from urllib.parse import urlparse, parse_qs
            params = parse_qs(urlparse(self.path).query)
            engine_name = params.get("engine", [""])[0]

            if engine_name not in ENGINES:
                self.send_response(400)
                self.send_header("Content-Type", "text/plain")
                self.end_headers()
                self.wfile.write(f"Unknown engine: {engine_name}".encode())
                return

            with engine_lock:
                if current_engine.name == engine_name:
                    self.send_response(200)
                    self.send_header("Content-Type", "text/plain")
                    self.end_headers()
                    self.wfile.write(b"already active")
                    return

                log(f"Switching engine: {current_engine.name} -> {engine_name}")
                new_engine = ENGINES[engine_name]()
                # Pre-load model in background to minimize gap
                try:
                    new_engine.load()
                except Exception as e:
                    self.send_response(500)
                    self.send_header("Content-Type", "text/plain")
                    self.end_headers()
                    self.wfile.write(f"Failed to load {engine_name}: {e}".encode())
                    log(f"Engine load failed: {e}")
                    return
                current_engine = new_engine

            broadcast({"type": "engine", "engine": engine_name})
            self.send_response(200)
            self.send_header("Content-Type", "text/plain")
            self.end_headers()
            self.wfile.write(f"switched to {engine_name}".encode())
            log(f"Engine switched to {engine_name}")

        elif self.path == "/engine":
            with engine_lock:
                name = current_engine.name if current_engine else "unknown"
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(json.dumps({"engine": name}).encode())

        elif self.path == "/stream":
            self.send_response(200)
            self.send_header("Content-Type", "text/event-stream")
            self.send_header("Cache-Control", "no-cache")
            self.send_header("Connection", "keep-alive")
            self.end_headers()

            q = queue.Queue()
            with sub_lock:
                subscribers.append(q)
            try:
                while running:
                    try:
                        data = q.get(timeout=15)
                        self.wfile.write(data)
                        self.wfile.flush()
                    except queue.Empty:
                        self.wfile.write(b": keepalive\n\n")
                        self.wfile.flush()
            except (BrokenPipeError, ConnectionResetError, OSError):
                pass
            finally:
                with sub_lock:
                    if q in subscribers:
                        subscribers.remove(q)

        else:
            self.send_error(404)

    def log_message(self, *_):
        pass


# ── Main ──────────────────────────────────────────────────────────────────────
def main():
    global running, current_engine

    import platform

    if platform.machine() != "arm64" or platform.system() != "Darwin":
        print(
            "TextStream requires Apple Silicon (M1/M2/M3/M4). "
            "MLX does not run on Intel Macs or Linux.",
            file=sys.stderr,
        )
        sys.exit(1)

    parser = argparse.ArgumentParser(description="TextStream - live speech to text")
    parser.add_argument("--port", type=int, default=DEFAULT_PORT)
    parser.add_argument(
        "--engine", default="qwen",
        choices=list(ENGINES.keys()),
        help="ASR engine (default: qwen)",
    )
    parser.add_argument(
        "--interval", type=float, default=DEFAULT_INTERVAL,
        help=f"Seconds between streaming updates (default: {DEFAULT_INTERVAL})",
    )
    parser.add_argument(
        "--vad-threshold", type=float, default=DEFAULT_VAD_THRESHOLD,
        help=f"Silero VAD speech probability threshold (default: {DEFAULT_VAD_THRESHOLD})",
    )
    parser.add_argument("--no-browser", action="store_true")
    parser.add_argument("--no-grafana", action="store_true", help="Disable Grafana push")
    args = parser.parse_args()

    if args.no_grafana or not GRAFANA_TOKEN:
        global push_annotation
        push_annotation = lambda text: None

    # Create and load initial engine
    current_engine = ENGINES[args.engine]()
    current_engine.load()

    stream = sd.InputStream(
        samplerate=SAMPLE_RATE,
        channels=1,
        dtype="float32",
        callback=audio_callback,
        blocksize=int(SAMPLE_RATE * 0.1),
    )
    stream.start()
    log("Microphone active")

    t = threading.Thread(
        target=transcription_loop, args=(args.interval, args.vad_threshold), daemon=True
    )
    t.start()

    httpd = ThreadedServer(("127.0.0.1", args.port), Handler)

    url = f"http://localhost:{args.port}"
    log(f"Browser:  {url}")
    log(f"Grafana:  annotations tagged 'textstream'")
    log(f"Interval: {args.interval}s streaming updates")

    if not args.no_browser:
        webbrowser.open(url)

    def shutdown(sig, frame):
        global running
        running = False
        log("\nStopping...")
        stream.stop()
        stream.close()
        httpd.shutdown()
        sys.exit(0)

    signal.signal(signal.SIGINT, shutdown)
    signal.signal(signal.SIGTERM, shutdown)

    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        shutdown(None, None)


if __name__ == "__main__":
    main()
