import {spawn, type ChildProcess} from 'node:child_process'
import * as fs from 'node:fs'
import * as path from 'node:path'
import {mergeTurns, renderJson, renderMarkdown, type Item} from './render.js'
import type {WavWriter} from './wav.js'

// Live transcription during recording (the record command's --transcribe).
//
// A single persistent python process loads faster-whisper ONCE and then
// serves transcription requests over a line-delimited JSON stdin/stdout
// protocol; each request names a WAV file and a frame range, which the helper
// reads raw from disk (only flushed bytes are ever requested, so reading a
// file that is still being appended to is safe — the header is bypassed).
//
// The recorder feeds every decoded audio frame through onAudio, which tracks
// silence so chunks can be cut at natural pauses (>= 300 ms below the RMS
// threshold) — never mid-word. Every tick, each segment with >= MIN_CHUNK of
// speech up to a silence cut is dispatched (or force-cut at MAX_CHUNK of
// unbroken speech). transcript.md / transcript.json are re-rendered after
// every result, so the transcript grows while the meeting is happening.
//
// If the helper dies, recording is NEVER affected: live transcription
// disables itself with a log line, and `transcribe` can be run afterwards.

// Tuned for ~8 s from end-of-utterance to text on screen: pause detect
// (~0.3 s) + WAV flush (~0.5 s avg) + tick (~1.5 s avg) + ASR (1-4 s with
// `small` on CPU) + render debounce (0.5 s). Smaller chunks trade a little
// per-chunk context/efficiency for latency; turn merging glues the text back.
const QUIET_RMS = 300 // ~ -41 dBFS: below this a 10 ms frame counts as quiet
const QUIET_CUT_MS = 300 // this much consecutive quiet = a safe cut point
const TICK_MS = 3000
const MIN_CHUNK_SEC = 2 // don't bother the model with less than this
const MAX_CHUNK_SEC = 30 // force a cut after this much unbroken speech
const FINISH_TIMEOUT_MS = 180000

interface AsrSegment {
  start: number
  end: number
  text: string
}

const LIVE_PY = `
import sys, json
import numpy as np
from faster_whisper import WhisperModel
model_name = sys.argv[1]
language = None if sys.argv[2] == "-" else sys.argv[2]
model = WhisperModel(model_name, device="auto", compute_type="auto")
print(json.dumps({"ready": True}), flush=True)
while True:
    # readline, not iteration: iterating sys.stdin read-ahead-buffers and can
    # sit on a complete line without yielding it
    line = sys.stdin.readline()
    if not line:
        break
    req = json.loads(line)
    ch = req["channels"]
    with open(req["wav"], "rb") as f:
        f.seek(44 + req["startFrame"] * ch * 2)
        raw = f.read((req["endFrame"] - req["startFrame"]) * ch * 2)
    a = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0
    if ch > 1:
        a = a.reshape(-1, ch).mean(axis=1)
    rate = req["rate"]
    if rate != 16000:
        if rate % 16000 == 0:
            k = rate // 16000
            n = (len(a) // k) * k
            a = a[:n].reshape(-1, k).mean(axis=1)
        else:
            xi = np.arange(0, len(a), rate / 16000.0)
            a = np.interp(xi, np.arange(len(a)), a)
    segs, info = model.transcribe(a.astype(np.float32), language=language, vad_filter=True)
    out = [{"start": s.start, "end": s.end, "text": s.text} for s in segs]
    print(json.dumps({"id": req["id"], "segments": out}), flush=True)
`

interface AsrRequest {
  wav: string
  rate: number
  channels: number
  startFrame: number
  endFrame: number
}

/** The persistent faster-whisper helper process. */
class PythonAsr {
  private proc: ChildProcess
  private pending = new Map<
    number,
    {resolve: (segs: AsrSegment[]) => void; reject: (err: Error) => void}
  >()
  private nextId = 1
  private buf = ''
  private readyResolve!: () => void
  private readyReject!: (err: Error) => void
  readonly ready: Promise<void>
  dead = false

  constructor(model: string, language: string | null, onLog: (l: string) => void) {
    this.ready = new Promise((resolve, reject) => {
      this.readyResolve = resolve
      this.readyReject = reject
    })
    // NB: `python3 -c <script>` — NOT `python3 -` + script on stdin, which
    // would read stdin to EOF before executing and never see our requests.
    // With -c, sys.argv is ['-c', model, language].
    this.proc = spawn('python3', ['-c', LIVE_PY, model, language ?? '-'], {
      stdio: ['pipe', 'pipe', 'pipe']
    })
    this.proc.stdout!.on('data', d => {
      this.buf += d
      const lines = this.buf.split('\n')
      this.buf = lines.pop() ?? ''
      for (const line of lines) {
        if (!line.trim()) continue
        let msg: {ready?: boolean; id?: number; segments?: AsrSegment[]}
        try {
          msg = JSON.parse(line)
        } catch {
          continue
        }
        if (msg.ready) {
          this.readyResolve()
          continue
        }
        if (typeof msg.id === 'number') {
          const p = this.pending.get(msg.id)
          this.pending.delete(msg.id)
          p?.resolve(msg.segments ?? [])
        }
      }
    })
    let errTail = ''
    this.proc.stderr!.on('data', d => {
      const line = String(d).trim()
      if (line) errTail = line.slice(0, 300)
    })
    this.proc.on('error', err => this.die(err.message))
    this.proc.on('exit', code => {
      if (!this.dead && code !== 0) {
        this.die(`helper exited with code ${code}${errTail ? ` (${errTail})` : ''}`)
        onLog(`live transcription disabled: ${errTail || `helper exited (${code})`}`)
      }
    })
  }

  private die(reason: string) {
    if (this.dead) return
    this.dead = true
    this.readyReject(new Error(reason))
    // ready may already be resolved; a handled rejection after that is fine
    this.ready.catch(() => undefined)
    for (const p of this.pending.values()) p.reject(new Error(reason))
    this.pending.clear()
  }

  request(req: AsrRequest): Promise<AsrSegment[]> {
    if (this.dead) return Promise.reject(new Error('helper is dead'))
    const id = this.nextId++
    return new Promise((resolve, reject) => {
      this.pending.set(id, {resolve, reject})
      this.proc.stdin!.write(JSON.stringify({id, ...req}) + '\n')
    })
  }

  close() {
    this.dead = true
    try {
      this.proc.stdin!.end()
    } catch {
      /* ignore */
    }
    setTimeout(() => {
      try {
        this.proc.kill('SIGKILL')
      } catch {
        /* ignore */
      }
    }, 2000).unref()
  }
}

interface LiveSeg {
  file: string // relative, as in the manifest
  absPath: string
  name: string
  rate: number
  channels: number
  startMs: number
  writer: WavWriter | null // null once the segment has ended
  /** Frames handed to ASR (or skipped as silence) so far. */
  transcribedUpTo: number
  /** Latest safe cut point (end of a >= 300 ms quiet stretch). */
  lastCutFrame: number
  lastEndFrame: number
  quietMs: number
  /** Any non-quiet audio since transcribedUpTo? */
  hasSpeech: boolean
  utterances: AsrSegment[] // in-file seconds
  ended: boolean
  endFrames: number
  pendingOps: number
  cacheWritten: boolean
}

export interface LiveTranscriberOptions {
  outDir: string
  room: string
  model: string | null
  language: string | null
  startedAtMs: number
  onLog: (line: string) => void
}

export class LiveTranscriber {
  private asr: PythonAsr
  private segs = new Map<string, LiveSeg>()
  private eventItems: Item[] = []
  private inFlight = new Set<Promise<unknown>>()
  private timer: ReturnType<typeof setInterval> | null = null
  private renderTimer: ReturnType<typeof setTimeout> | null = null
  private model: string

  constructor(private opts: LiveTranscriberOptions) {
    this.model = opts.model ?? 'small'
    this.asr = new PythonAsr(this.model, opts.language, opts.onLog)
    opts.onLog(`live transcription: loading model ${this.model}...`)
    void this.asr.ready.then(
      () => opts.onLog('live transcription ready'),
      () => undefined // logged by PythonAsr
    )
    this.timer = setInterval(() => this.tick(), TICK_MS)
  }

  onSegmentStart(
    file: string,
    absPath: string,
    name: string,
    rate: number,
    channels: number,
    startMs: number,
    writer: WavWriter
  ) {
    this.segs.set(file, {
      file,
      absPath,
      name,
      rate,
      channels,
      startMs,
      writer,
      transcribedUpTo: 0,
      lastCutFrame: 0,
      lastEndFrame: 0,
      quietMs: 0,
      hasSpeech: false,
      utterances: [],
      ended: false,
      endFrames: 0,
      pendingOps: 0,
      cacheWritten: false
    })
  }

  /** Called for every decoded frame batch; endFrame = writer.framesWritten
   *  AFTER appending (so padded gaps show up as position jumps). */
  onAudio(file: string, endFrame: number, samples: Int16Array) {
    const seg = this.segs.get(file)
    if (!seg || seg.ended) return
    const frames = samples.length / seg.channels
    const startFrame = endFrame - frames
    if (startFrame > seg.lastEndFrame) {
      // A silence-padded gap was inserted before this batch: it is all quiet,
      // so the start of the current batch is a safe cut point.
      seg.quietMs = 0
      seg.lastCutFrame = startFrame
    }
    seg.lastEndFrame = endFrame
    let sumSq = 0
    for (let i = 0; i < samples.length; i++) sumSq += samples[i]! * samples[i]!
    const rms = Math.sqrt(sumSq / samples.length)
    if (rms < QUIET_RMS) {
      seg.quietMs += (frames / seg.rate) * 1000
      if (seg.quietMs >= QUIET_CUT_MS) seg.lastCutFrame = endFrame
    } else {
      seg.quietMs = 0
      seg.hasSpeech = true
    }
  }

  onSegmentEnd(file: string) {
    const seg = this.segs.get(file)
    if (!seg || seg.ended) return
    seg.ended = true
    seg.endFrames = seg.writer?.framesWritten ?? seg.lastEndFrame
    seg.writer = null
    this.dispatch(seg, seg.endFrames)
    this.maybeWriteCache(seg)
  }

  /** Chat / join / left items from the recorder, for the rendered timeline. */
  onEvent(item: Item) {
    this.eventItems.push(item)
    this.scheduleRender()
  }

  private tick() {
    for (const seg of this.segs.values()) {
      if (seg.ended || !seg.writer) continue
      const flushed = seg.writer.flushedFrames
      const cut = Math.min(seg.lastCutFrame, flushed)
      if (cut - seg.transcribedUpTo >= MIN_CHUNK_SEC * seg.rate) {
        this.dispatch(seg, cut)
      } else if (flushed - seg.transcribedUpTo >= MAX_CHUNK_SEC * seg.rate) {
        this.dispatch(seg, flushed) // unbroken speech: cut anyway
      }
    }
  }

  private dispatch(seg: LiveSeg, to: number) {
    const from = seg.transcribedUpTo
    if (to <= from) return
    seg.transcribedUpTo = to
    const hadSpeech = seg.hasSpeech
    seg.hasSpeech = false
    if (!hadSpeech || this.asr.dead) return
    seg.pendingOps++
    const op = this.asr
      .request({
        wav: seg.absPath,
        rate: seg.rate,
        channels: seg.channels,
        startFrame: from,
        endFrame: to
      })
      .then(segments => {
        for (const s of segments) {
          const text = s.text.trim()
          if (!text) continue
          seg.utterances.push({
            start: s.start + from / seg.rate,
            end: s.end + from / seg.rate,
            text
          })
        }
        if (segments.length > 0) this.scheduleRender()
      })
      .catch(() => undefined) // helper death is logged once by PythonAsr
      .finally(() => {
        seg.pendingOps--
        this.inFlight.delete(op)
        this.maybeWriteCache(seg)
      })
    this.inFlight.add(op)
  }

  /** Once a segment has ended and drained, persist its ASR results in the
   *  same asr/*.json format the offline transcribe subcommand uses/caches. */
  private maybeWriteCache(seg: LiveSeg) {
    if (!seg.ended || seg.pendingOps > 0 || seg.cacheWritten) return
    seg.cacheWritten = true
    if (this.asr.dead && seg.utterances.length === 0) return // let offline redo it
    const asrDir = path.join(this.opts.outDir, 'asr')
    try {
      fs.mkdirSync(asrDir, {recursive: true})
      fs.writeFileSync(
        path.join(asrDir, path.basename(seg.file).replace(/\.wav$/i, '') + '.json'),
        JSON.stringify(
          {
            engine: 'faster-whisper',
            model: this.model,
            segments: [...seg.utterances].sort((a, b) => a.start - b.start)
          },
          null,
          2
        ) + '\n'
      )
    } catch {
      /* ignore */
    }
  }

  private scheduleRender() {
    if (this.renderTimer !== null) return
    this.renderTimer = setTimeout(() => {
      this.renderTimer = null
      this.render(null)
    }, 500)
  }

  private render(endedAtMs: number | null) {
    const items: Item[] = [...this.eventItems]
    for (const seg of this.segs.values()) {
      for (const u of seg.utterances) {
        items.push({
          timeMs: seg.startMs + u.start * 1000,
          endMs: seg.startMs + u.end * 1000,
          type: 'speech',
          speaker: seg.name,
          text: u.text
        })
      }
    }
    const merged = mergeTurns(items)
    const meta = {
      room: this.opts.room,
      startedAtMs: this.opts.startedAtMs,
      endedAtMs,
      speakers: [...new Set([...this.segs.values()].map(s => s.name))],
      engine: 'faster-whisper (live)',
      model: this.model
    }
    try {
      fs.writeFileSync(
        path.join(this.opts.outDir, 'transcript.md'),
        renderMarkdown(meta, merged)
      )
      fs.writeFileSync(
        path.join(this.opts.outDir, 'transcript.json'),
        renderJson(meta, merged)
      )
    } catch {
      /* ignore */
    }
  }

  /** All segments have been ended by the recorder; drain and finalize. */
  async finish(endedAtMs: number): Promise<void> {
    if (this.timer !== null) clearInterval(this.timer)
    this.timer = null
    if (this.renderTimer !== null) clearTimeout(this.renderTimer)
    this.renderTimer = null
    if (this.inFlight.size > 0) {
      this.opts.onLog(
        `waiting for ${this.inFlight.size} transcription chunk(s) to finish...`
      )
      await Promise.race([
        Promise.allSettled([...this.inFlight]),
        new Promise(r => setTimeout(r, FINISH_TIMEOUT_MS))
      ])
    }
    this.asr.close()
    for (const seg of this.segs.values()) this.maybeWriteCache(seg)
    this.render(endedAtMs)
    this.opts.onLog(`wrote ${path.join(this.opts.outDir, 'transcript.md')}`)
  }
}