#![allow(clippy::too_many_arguments)]

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{
    atomic::{AtomicBool, AtomicU64, Ordering},
    Arc, Mutex,
};
use std::thread;
use std::time::{Duration, Instant};

use serde::{Deserialize, Serialize};
use tauri::ipc::Channel;

use crate::commands::audio::resolve_recording_path;
use crate::commands::build_initial_prompt;
use crate::commands::models::{default_model_id_for_engine, ensure_model_loaded};
use crate::commands::power::PowerAssertion;
use crate::AppState;
use kon_ai_formatting::{post_process_segments, FormatMode, PostProcessOptions};
use kon_audio::{MicrophoneCapture, StreamingResampler, WavWriter};
use kon_core::constants::WHISPER_SAMPLE_RATE;
use kon_core::types::{AudioSamples, Segment, TranscriptionOptions};
use kon_transcription::LocalEngine;

const CHUNK_SAMPLES: usize = 32_000; // 2s at 16kHz
const OVERLAP_SAMPLES: usize = 4_000; // 0.25s at 16kHz
const FINAL_CHUNK_MIN_SAMPLES: usize = 4_000; // 0.25s
const MAX_PENDING_SAMPLES: usize = CHUNK_SAMPLES;
const SPEECH_FRAME_SAMPLES: usize = 800; // 50ms
const MIN_SPEECH_FRAMES: usize = 1; // any plausible speech-like frame
const SILENCE_RMS_THRESHOLD: f32 = 0.001;
const SPEECH_WINDOW_RMS_THRESHOLD: f32 = 0.0014;
const SPEECH_WINDOW_PEAK_THRESHOLD: f32 = 0.004;
const STRONG_SPEECH_RMS_THRESHOLD: f32 = 0.003;
const STRONG_SPEECH_PEAK_THRESHOLD: f32 = 0.012;
const FLATLINE_PEAK_THRESHOLD: f32 = 0.0005;
const DUPLICATE_TRANSCRIPT_WINDOW_SECS: f64 = 6.0;
const DUPLICATE_TRANSCRIPT_MERGE_LIMIT: usize = 3;
const DUPLICATE_HISTORY_RETENTION_SECS: f64 = 8.0;
const DUPLICATE_CHECK_LEADING_SECS: f64 = 1.5;
const TOKEN_COVERAGE_THRESHOLD: f64 = 0.6;
const TOKEN_SEQUENCE_THRESHOLD: f64 = 0.6;
const MIN_TOKENS_FOR_OVERLAP: usize = 3;
const MEANINGFUL_TOKEN_COVERAGE_THRESHOLD: f64 = 0.55;
const MEANINGFUL_TOKEN_SEQUENCE_THRESHOLD: f64 = 0.55;
const MIN_MEANINGFUL_TOKENS_FOR_OVERLAP: usize = 4;
const LOW_SIGNAL_TOKENS: &[&str] = &[
    "a", "an", "and", "are", "as", "at", "be", "been", "being", "but", "by", "d", "did", "do",
    "does", "for", "from", "had", "has", "have", "he", "her", "here", "his", "how", "i", "if",
    "in", "is", "it", "ll", "m", "me", "my", "of", "on", "or", "our", "out", "re", "s", "she",
    "so", "t", "that", "the", "their", "them", "there", "these", "they", "this", "those", "to",
    "ve", "was", "we", "well", "were", "what", "when", "where", "which", "who", "why", "with",
    "without", "you", "your",
];

#[derive(Default)]
pub struct LiveTranscriptionState {
    next_session_id: AtomicU64,
    running: Mutex<Option<RunningLiveSession>>,
}

struct RunningLiveSession {
    id: u64,
    stop_flag: Arc<AtomicBool>,
    handle: tokio::task::JoinHandle<Result<LiveSessionSummary, String>>,
    status_channel: Channel<LiveStatusMessage>,
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct StartLiveTranscriptionConfig {
    pub engine: String,
    pub model_id: Option<String>,
    pub language: Option<String>,
    pub initial_prompt: Option<String>,
    pub save_audio: bool,
    pub output_folder: Option<String>,
    pub remove_fillers: bool,
    pub british_english: bool,
    pub anti_hallucination: bool,
    pub format_mode: String,
    /// Optional explicit microphone device name (from `list_audio_devices`).
    /// None or empty string = let `MicrophoneCapture::start` auto-select.
    pub microphone_device: Option<String>,
    /// Optional profile id. None falls back to `DEFAULT_PROFILE_ID`. Drives
    /// the post-processing dictionary via `profile_terms` and, when the
    /// request's `initial_prompt` is empty, supplies the Whisper prompt.
    pub profile_id: Option<String>,
}

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct StartLiveTranscriptionResponse {
    pub session_id: u64,
}

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct StopLiveTranscriptionResponse {
    pub session_id: u64,
    pub audio_path: Option<String>,
    pub dropped_audio_ms: u64,
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct LiveResultMessage {
    pub session_id: u64,
    pub chunk_id: u32,
    pub chunk_start_secs: f64,
    pub duration: f64,
    pub language: String,
    pub inference_ms: u64,
    pub segments: Vec<Segment>,
    /// Concatenated text BEFORE post-processing (no filler removal, no
    /// British conversion, no LLM cleanup). Used by the transcription
    /// preview overlay so the user can see raw Whisper output as it
    /// streams in.
    pub raw_text: String,
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase", tag = "type")]
#[allow(dead_code)]
pub enum LiveStatusMessage {
    Warning {
        session_id: u64,
        message: String,
    },
    Overload {
        session_id: u64,
        dropped_audio_ms: u64,
        message: String,
    },
    Error {
        session_id: u64,
        message: String,
    },
    Finished {
        session_id: u64,
        audio_path: Option<String>,
        dropped_audio_ms: u64,
    },
}

struct LiveSessionSummary {
    session_id: u64,
    dropped_audio_ms: u64,
    /// Absolute path of the progressively-written WAV file. `Some` iff
    /// `save_audio` was requested and the writer opened successfully;
    /// the file on disk is a valid, playable WAV even if the session
    /// was killed mid-append (brief item #19 crash-safety guarantee).
    audio_path: Option<String>,
}

struct InferenceTask {
    chunk_id: u32,
    chunk_start_sample: u64,
    trim_before_secs: f64,
    duration_secs: f64,
    rx: std::sync::mpsc::Receiver<Result<kon_transcription::TimedTranscript, String>>,
}

#[derive(Debug, Clone)]
struct RecentTranscriptSegment {
    start_secs: f64,
    end_secs: f64,
    text: String,
}

#[derive(Debug, Clone, Copy, Default)]
struct SpeechGateState {
    peak_rms: f32,
    peak_amplitude: f32,
    window_count: usize,
    speech_window_count: usize,
    consecutive_speech_windows: usize,
    max_consecutive_speech_windows: usize,
}

#[derive(Debug, Clone, Copy, PartialEq)]
struct SpeechGateDecision {
    skip: bool,
    reason: &'static str,
    peak_rms: f32,
    peak_amplitude: f32,
    window_count: usize,
    speech_window_count: usize,
    max_consecutive_speech_windows: usize,
}

#[tauri::command]
pub async fn start_live_transcription_session(
    app: tauri::AppHandle,
    state: tauri::State<'_, AppState>,
    live_state: tauri::State<'_, LiveTranscriptionState>,
    mut config: StartLiveTranscriptionConfig,
    result_channel: Channel<LiveResultMessage>,
    status_channel: Channel<LiveStatusMessage>,
) -> Result<StartLiveTranscriptionResponse, String> {
    {
        let running = live_state.running.lock().unwrap();
        if running.is_some() {
            return Err("A live transcription session is already running".into());
        }
    }

    let resolved_profile_id = config
        .profile_id
        .clone()
        .unwrap_or_else(|| kon_storage::DEFAULT_PROFILE_ID.to_string());

    let profile = kon_storage::database::get_profile(&state.db, &resolved_profile_id)
        .await
        .map_err(|e| e.to_string())?
        .ok_or_else(|| format!("Profile {resolved_profile_id} not found"))?;

    let profile_terms: Vec<String> =
        kon_storage::database::list_profile_terms(&state.db, &resolved_profile_id)
            .await
            .map_err(|e| e.to_string())?
            .into_iter()
            .map(|t| t.term)
            .collect();

    // Collapse the effective initial_prompt on the struct so downstream
    // `TranscriptionOptions` construction (see `maybe_dispatch_chunk`) picks
    // up profile fallback + vocabulary injection without further plumbing.
    let request_prompt = config.initial_prompt.clone().unwrap_or_default();
    config.initial_prompt = build_initial_prompt(
        &request_prompt,
        &profile.initial_prompt,
        &profile_terms,
    );

    let model_id = config
        .model_id
        .clone()
        .unwrap_or_else(|| default_model_id_for_engine(&config.engine).to_string());
    eprintln!(
        "[live] starting session: engine={}, model={}, language={:?}, save_audio={}",
        config.engine, model_id, config.language, config.save_audio
    );
    // None: live-transcription model loads don't enforce sequential-GPU
    // mode. The Settings-level load flow owns that guard.
    ensure_model_loaded(&state, &config.engine, &model_id, None).await?;

    let session_id = live_state
        .next_session_id
        .fetch_add(1, Ordering::Relaxed)
        .saturating_add(1);
    let stop_flag = Arc::new(AtomicBool::new(false));
    let engine = pick_engine(&state, &config.engine)?;

    // Resolve the WAV destination up front so the progressive writer
    // (brief item #19) can open it before any samples arrive. Failure
    // to create the recordings directory is fatal — the user asked
    // for save_audio=true and silently dropping the recording would
    // surprise them worse.
    let audio_path = if config.save_audio {
        Some(resolve_recording_path(&app, config.output_folder.as_deref())?)
    } else {
        None
    };

    let worker_stop = stop_flag.clone();
    let worker_status = status_channel.clone();
    let worker_results = result_channel.clone();

    let dictionary_terms = profile_terms.clone();
    let worker_audio_path = audio_path.clone();

    let handle = tokio::task::spawn_blocking(move || {
        run_live_session(
            session_id,
            engine,
            config,
            worker_audio_path,
            dictionary_terms,
            worker_results,
            worker_status,
            worker_stop,
        )
    });

    *live_state.running.lock().unwrap() = Some(RunningLiveSession {
        id: session_id,
        stop_flag,
        handle,
        status_channel,
    });

    Ok(StartLiveTranscriptionResponse { session_id })
}

#[tauri::command]
pub async fn stop_live_transcription_session(
    app: tauri::AppHandle,
    live_state: tauri::State<'_, LiveTranscriptionState>,
    session_id: u64,
) -> Result<StopLiveTranscriptionResponse, String> {
    let running = live_state.running.lock().unwrap().take();
    let Some(running) = running else {
        return Err("No live transcription session is running".into());
    };

    if running.id != session_id {
        *live_state.running.lock().unwrap() = Some(running);
        return Err(format!("Session {session_id} is not active"));
    }

    running.stop_flag.store(true, Ordering::Relaxed);

    let summary = running
        .handle
        .await
        .map_err(|e| format!("Live session task failed: {e}"))??;

    // Progressive WAV writer (brief item #19) wrote samples to disk
    // throughout the session; the path is already finalised. Nothing
    // further to persist.
    let _ = app;
    let audio_path = summary.audio_path;

    let response = StopLiveTranscriptionResponse {
        session_id: summary.session_id,
        audio_path: audio_path.clone(),
        dropped_audio_ms: summary.dropped_audio_ms,
    };

    let _ = running.status_channel.send(LiveStatusMessage::Finished {
        session_id: summary.session_id,
        audio_path,
        dropped_audio_ms: summary.dropped_audio_ms,
    });

    Ok(response)
}

fn pick_engine(state: &AppState, engine: &str) -> Result<Arc<LocalEngine>, String> {
    match engine {
        "whisper" => Ok(state.whisper_engine.clone()),
        "parakeet" => Ok(state.parakeet_engine.clone()),
        other => Err(format!("Unknown engine: {other}")),
    }
}

fn run_live_session(
    session_id: u64,
    engine: Arc<LocalEngine>,
    config: StartLiveTranscriptionConfig,
    audio_path: Option<PathBuf>,
    dictionary_terms: Vec<String>,
    result_channel: Channel<LiveResultMessage>,
    status_channel: Channel<LiveStatusMessage>,
    stop_flag: Arc<AtomicBool>,
) -> Result<LiveSessionSummary, String> {
    // macOS: disable App Nap while recording. On every other OS this
    // is a no-op. Keeping the guard in scope ties the assertion's
    // lifetime to the session — when the function returns, the Drop
    // impl lifts it. Item #9 in docs/whisper-ecosystem/brief.md.
    let _power_guard = PowerAssertion::begin("kon live dictation session");

    let (mut capture, rx) = match config.microphone_device.as_deref() {
        Some(name) if !name.is_empty() => MicrophoneCapture::start_with_device(name),
        _ => MicrophoneCapture::start(),
    }
    .map_err(|e| e.to_string())?;
    // Drain runtime stream errors into the status channel so the user
    // gets a toast when the device disconnects mid-recording instead of
    // silently producing empty transcripts. The `_capture` binding keeps
    // the cpal stream alive for the duration of the session.
    let mic_error_rx = capture.take_error_rx();
    let _capture = capture;

    let mut resampler: Option<StreamingResampler> = None;
    let mut capture_buffer: Vec<f32> = Vec::new();

    // Progressive WAV writer (brief item #19). Sample rate comes from
    // the loaded backend's capabilities (#13 wiring) so a future
    // non-16kHz backend records at its native rate without further
    // plumbing. The writer flushes its header every ~500 ms, so the
    // file on disk is a playable WAV even if the process is killed.
    let sample_rate = engine
        .capabilities()
        .map(|c| c.sample_rate)
        .unwrap_or(WHISPER_SAMPLE_RATE);
    let mut wav_writer: Option<WavWriter> = match audio_path.as_ref() {
        Some(path) => match WavWriter::create(path, sample_rate, 1) {
            Ok(w) => Some(w),
            Err(e) => {
                let _ = status_channel.send(LiveStatusMessage::Warning {
                    session_id,
                    message: format!(
                        "Failed to open audio recording file ({}); transcription will continue without saving audio.",
                        e
                    ),
                });
                None
            }
        },
        None => None,
    };
    // `reported_audio_path` is decided at end-of-session based on
    // whether the writer finalised successfully, not at open time.
    // This way a writer that dies mid-session (append error clearing
    // wav_writer, or finalise returning Err) does not leak a stale
    // path back to the frontend that might point to a file whose
    // header is out of sync with its data chunk.

    let mut buffer_start_sample: u64 = 0;
    let mut dropped_audio_ms: u64 = 0;
    let mut chunk_id: u32 = 0;
    let mut inflight: Option<InferenceTask> = None;
    let mut resampler_flushed = false;
    let mut recent_segments: Vec<RecentTranscriptSegment> = Vec::new();

    loop {
        if let Some(_done) = poll_inference(
            &mut inflight,
            session_id,
            &config,
            &mut recent_segments,
            &dictionary_terms,
            &result_channel,
            &status_channel,
        )? {}

        // Surface any cpal runtime errors as warnings. Non-fatal: a hard
        // disconnect will also drop the audio sender and be caught by
        // the `Disconnected` arm below. This lets the user see a toast
        // even when cpal recovers without tearing the stream down.
        if let Some(err_rx) = &mic_error_rx {
            while let Ok(err) = err_rx.try_recv() {
                let _ = status_channel.send(LiveStatusMessage::Warning {
                    session_id,
                    message: format!(
                        "Microphone '{}' reported an error: {}",
                        err.device_name, err.message
                    ),
                });
            }
        }

        match rx.recv_timeout(Duration::from_millis(25)) {
            Ok(chunk) => {
                let mono = downmix_chunk(chunk.samples, chunk.channels as usize);
                let resampler = match &mut resampler {
                    Some(resampler) => resampler,
                    None => {
                        resampler = Some(
                            StreamingResampler::new(chunk.sample_rate)
                                .map_err(|e| e.to_string())?,
                        );
                        resampler.as_mut().expect("resampler just set")
                    }
                };

                let resampled = resampler.push_samples(&mono).map_err(|e| e.to_string())?;
                append_resampled_audio(
                    &mut capture_buffer,
                    &mut wav_writer,
                    &resampled,
                    session_id,
                    &status_channel,
                );
            }
            Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {}
            Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => {
                let message = "Microphone capture disconnected unexpectedly".to_string();
                let _ = status_channel.send(LiveStatusMessage::Error {
                    session_id,
                    message: message.clone(),
                });
                return Err(message);
            }
        }

        if inflight.is_some() && capture_buffer.len() > MAX_PENDING_SAMPLES {
            let overflow = capture_buffer.len() - MAX_PENDING_SAMPLES;
            capture_buffer.drain(..overflow);
            buffer_start_sample = buffer_start_sample.saturating_add(overflow as u64);
            dropped_audio_ms = dropped_audio_ms
                .saturating_add((overflow as u64 * 1000) / WHISPER_SAMPLE_RATE as u64);
            let _ = status_channel.send(LiveStatusMessage::Overload {
                session_id,
                dropped_audio_ms,
                message: "Kon dropped older audio to keep live dictation responsive".into(),
            });
        }

        let stopping = stop_flag.load(Ordering::Relaxed);
        if stopping && !resampler_flushed {
            if let Some(resampler) = &mut resampler {
                let tail = resampler.flush().map_err(|e| e.to_string())?;
                append_resampled_audio(
                    &mut capture_buffer,
                    &mut wav_writer,
                    &tail,
                    session_id,
                    &status_channel,
                );
            }
            resampler_flushed = true;
            // Final flush for the WAV header so the last chunk's header
            // update is on disk before we drop into the inference drain.
            if let Some(writer) = wav_writer.as_mut() {
                if let Err(e) = writer.flush() {
                    let _ = status_channel.send(LiveStatusMessage::Warning {
                        session_id,
                        message: format!("WAV flush failed near session end: {e}"),
                    });
                }
            }
        }

        if inflight.is_none() {
            if let Some(task) = maybe_dispatch_chunk(
                &engine,
                &config,
                &mut capture_buffer,
                &mut buffer_start_sample,
                &mut chunk_id,
                stopping,
                &status_channel,
                session_id,
            ) {
                inflight = Some(task);
                continue;
            }

            if stopping && resampler_flushed {
                break;
            }
        }
    }

    while inflight.is_some() {
        poll_inference(
            &mut inflight,
            session_id,
            &config,
            &mut recent_segments,
            &dictionary_terms,
            &result_channel,
            &status_channel,
        )?;
        thread::sleep(Duration::from_millis(10));
    }

    // Finalise the progressive WAV writer and decide whether to
    // report a path to the frontend. Only a clean finalise produces a
    // reported path: a writer that died mid-session (wav_writer was
    // already None) or a finalise that itself errored both yield
    // `None`, so `StopLiveTranscriptionResponse.audio_path` reflects
    // "recording is known-good" rather than "recording was attempted".
    let audio_path = match wav_writer.take() {
        Some(writer) => match writer.finalize() {
            Ok(()) => audio_path.as_ref().map(|p| p.to_string_lossy().to_string()),
            Err(e) => {
                let _ = status_channel.send(LiveStatusMessage::Warning {
                    session_id,
                    message: format!(
                        "WAV finalise failed: {e}. A partial recording may still be on disk from earlier flushes."
                    ),
                });
                None
            }
        },
        None => None,
    };

    Ok(LiveSessionSummary {
        session_id,
        dropped_audio_ms,
        audio_path,
    })
}

fn append_resampled_audio(
    capture_buffer: &mut Vec<f32>,
    wav_writer: &mut Option<WavWriter>,
    resampled: &[f32],
    session_id: u64,
    status_channel: &Channel<LiveStatusMessage>,
) {
    if resampled.is_empty() {
        return;
    }

    capture_buffer.extend_from_slice(resampled);
    if let Some(writer) = wav_writer.as_mut() {
        if let Err(e) = writer.append(resampled) {
            // WAV write failure is non-fatal for live transcription —
            // drop the writer so subsequent samples don't keep trying
            // a broken file handle, and warn the user. The samples
            // already written up to this point remain playable thanks
            // to periodic header flushes.
            let _ = status_channel.send(LiveStatusMessage::Warning {
                session_id,
                message: format!(
                    "Audio recording halted: {e}. Transcription continues; partial WAV is playable."
                ),
            });
            *wav_writer = None;
        }
    }
}

fn maybe_dispatch_chunk(
    engine: &Arc<LocalEngine>,
    config: &StartLiveTranscriptionConfig,
    capture_buffer: &mut Vec<f32>,
    buffer_start_sample: &mut u64,
    chunk_id: &mut u32,
    stopping: bool,
    status_channel: &Channel<LiveStatusMessage>,
    session_id: u64,
) -> Option<InferenceTask> {
    let target_len = if capture_buffer.len() >= CHUNK_SAMPLES {
        CHUNK_SAMPLES
    } else if stopping && capture_buffer.len() >= FINAL_CHUNK_MIN_SAMPLES {
        capture_buffer.len()
    } else {
        return None;
    };

    let trim_before_secs = if *chunk_id > 0 && !stopping && target_len > OVERLAP_SAMPLES {
        OVERLAP_SAMPLES as f64 / WHISPER_SAMPLE_RATE as f64
    } else {
        0.0
    };

    let speech_window = if trim_before_secs > 0.0 {
        &capture_buffer[OVERLAP_SAMPLES..target_len]
    } else {
        &capture_buffer[..target_len]
    };

    let speech_gate = evaluate_speech_gate(speech_window);
    if speech_gate.skip {
        let skipped_ms = (target_len as u64 * 1000) / WHISPER_SAMPLE_RATE as u64;
        let gate_reason = match speech_gate.reason {
            "silence" => "near-silence",
            "insufficient_speech" => "insufficient speech energy",
            other => other,
        };
        eprintln!(
            "[live] session {session_id}: skipped {skipped_ms}ms chunk as {gate_reason} \
(peak_rms={:.6}, peak={:.6}, speech_windows={}/{}, max_consecutive={})",
            speech_gate.peak_rms,
            speech_gate.peak_amplitude,
            speech_gate.speech_window_count,
            speech_gate.window_count,
            speech_gate.max_consecutive_speech_windows,
        );
        let _ = status_channel.send(LiveStatusMessage::Warning {
            session_id,
            message: match speech_gate.reason {
                "silence" => format!(
                    "Skipped {skipped_ms}ms of near-silent audio. If this keeps happening, try a louder mic level or move closer to the microphone."
                ),
                _ => format!(
                    "Skipped {skipped_ms}ms of low-confidence audio. If this keeps happening, try a louder mic level or reduce background noise."
                ),
            },
        });
        capture_buffer.drain(..target_len);
        *buffer_start_sample = buffer_start_sample.saturating_add(target_len as u64);
        return None;
    }

    *chunk_id = chunk_id.saturating_add(1);
    let current_chunk_id = *chunk_id;
    let chunk_start_sample = *buffer_start_sample;
    let duration_secs = target_len as f64 / WHISPER_SAMPLE_RATE as f64;
    let chunk_samples = capture_buffer[..target_len].to_vec();
    eprintln!(
        "[live] session {session_id}: dispatching chunk {} ({duration_secs:.2}s, {} samples)",
        current_chunk_id,
        chunk_samples.len()
    );
    let advance_by = if stopping || target_len < CHUNK_SAMPLES {
        target_len
    } else {
        target_len.saturating_sub(OVERLAP_SAMPLES)
    };
    capture_buffer.drain(..advance_by);
    *buffer_start_sample = buffer_start_sample.saturating_add(advance_by as u64);

    let options = TranscriptionOptions {
        language: config.language.clone(),
        initial_prompt: config.initial_prompt.clone(),
    };
    let engine = engine.clone();
    let (tx, rx) = std::sync::mpsc::channel();

    thread::spawn(move || {
        let audio = AudioSamples::mono_16khz(chunk_samples);
        let started = Instant::now();
        let result = engine
            .transcribe_sync(&audio, &options)
            .map(|mut timed| {
                timed.inference_ms = started.elapsed().as_millis() as u64;
                timed
            })
            .map_err(|e| e.to_string());
        let _ = tx.send(result);
    });

    Some(InferenceTask {
        chunk_id: current_chunk_id,
        chunk_start_sample,
        trim_before_secs,
        duration_secs,
        rx,
    })
}

fn poll_inference(
    inflight: &mut Option<InferenceTask>,
    session_id: u64,
    config: &StartLiveTranscriptionConfig,
    recent_segments: &mut Vec<RecentTranscriptSegment>,
    dictionary_terms: &[String],
    result_channel: &Channel<LiveResultMessage>,
    status_channel: &Channel<LiveStatusMessage>,
) -> Result<Option<bool>, String> {
    let Some(task) = inflight else {
        return Ok(None);
    };

    match task.rx.try_recv() {
        Ok(Ok(timed)) => {
            let mut segments: Vec<Segment> = timed.transcript.segments().to_vec();
            trim_overlap_segments(&mut segments, task.trim_before_secs);
            // Capture raw text BEFORE any post-processing so the preview
            // overlay can show what Whisper actually returned.
            let raw_text = segments
                .iter()
                .map(|segment| segment.text.trim())
                .filter(|segment| !segment.is_empty())
                .collect::<Vec<_>>()
                .join(" ");
            post_process_segments(
                &mut segments,
                &PostProcessOptions {
                    remove_fillers: config.remove_fillers,
                    british_english: config.british_english,
                    anti_hallucination: config.anti_hallucination,
                    format_mode: FormatMode::parse(&config.format_mode),
                    dictionary_terms: dictionary_terms.to_vec(),
                },
                None,
            );
            let chunk_start_secs = task.chunk_start_sample as f64 / WHISPER_SAMPLE_RATE as f64;
            let skipped_duplicates = filter_duplicate_boundary_segments(
                &mut segments,
                chunk_start_secs,
                recent_segments,
            );
            let segment_count = segments.len();
            let delivered_segments = segments.clone();

            result_channel
                .send(LiveResultMessage {
                    session_id,
                    chunk_id: task.chunk_id,
                    chunk_start_secs,
                    duration: task.duration_secs,
                    language: timed.transcript.language().to_string(),
                    inference_ms: timed.inference_ms,
                    segments,
                    raw_text,
                })
                .map_err(|e| e.to_string())?;
            remember_recent_segments(recent_segments, &delivered_segments, chunk_start_secs);
            eprintln!(
                "[live] session {session_id}: delivered chunk {} with {} segments in {}ms{}",
                task.chunk_id,
                segment_count,
                timed.inference_ms,
                if skipped_duplicates > 0 {
                    format!(" (skipped {skipped_duplicates} duplicate boundary segment(s))")
                } else {
                    String::new()
                }
            );

            *inflight = None;
            Ok(Some(true))
        }
        Ok(Err(err)) => {
            eprintln!("[live] session {session_id}: inference error: {err}");
            *inflight = None;
            let _ = status_channel.send(LiveStatusMessage::Error {
                session_id,
                message: err.clone(),
            });
            Err(err)
        }
        Err(std::sync::mpsc::TryRecvError::Empty) => Ok(Some(false)),
        Err(std::sync::mpsc::TryRecvError::Disconnected) => {
            *inflight = None;
            let message = "Inference worker disconnected unexpectedly".to_string();
            eprintln!("[live] session {session_id}: {message}");
            let _ = status_channel.send(LiveStatusMessage::Error {
                session_id,
                message: message.clone(),
            });
            Err(message)
        }
    }
}

fn trim_overlap_segments(segments: &mut Vec<Segment>, trim_before_secs: f64) {
    if trim_before_secs <= 0.0 {
        return;
    }

    segments.retain(|segment| segment.end > trim_before_secs);
    for segment in segments.iter_mut() {
        if segment.start < trim_before_secs {
            segment.start = trim_before_secs;
        }
    }
}

fn filter_duplicate_boundary_segments(
    segments: &mut Vec<Segment>,
    chunk_start_secs: f64,
    recent_segments: &[RecentTranscriptSegment],
) -> usize {
    if recent_segments.is_empty() {
        return 0;
    }

    let mut skipped = 0usize;
    segments.retain(|segment| {
        if segment.start > DUPLICATE_CHECK_LEADING_SECS {
            return true;
        }

        let absolute_start = chunk_start_secs + segment.start;
        let candidates = build_nearby_transcript_candidates(recent_segments, absolute_start);
        if candidates.is_empty() {
            return true;
        }

        let duplicate = candidates.iter().any(|candidate| {
            transcripts_overlap(&segment.text, candidate)
                || transcripts_loosely_overlap(&segment.text, candidate)
        });

        if duplicate {
            skipped += 1;
            return false;
        }

        true
    });

    skipped
}

fn remember_recent_segments(
    recent_segments: &mut Vec<RecentTranscriptSegment>,
    segments: &[Segment],
    chunk_start_secs: f64,
) {
    if segments.is_empty() {
        return;
    }

    for segment in segments {
        let text = segment.text.trim();
        if text.is_empty() {
            continue;
        }

        recent_segments.push(RecentTranscriptSegment {
            start_secs: chunk_start_secs + segment.start,
            end_secs: chunk_start_secs + segment.end,
            text: text.to_string(),
        });
    }

    let cutoff = recent_segments
        .last()
        .map(|segment| segment.end_secs - DUPLICATE_HISTORY_RETENTION_SECS)
        .unwrap_or(0.0);
    recent_segments.retain(|segment| segment.end_secs >= cutoff);
}

fn build_nearby_transcript_candidates(
    recent_segments: &[RecentTranscriptSegment],
    timestamp_secs: f64,
) -> Vec<String> {
    let mut nearby: Vec<&RecentTranscriptSegment> = recent_segments
        .iter()
        .filter(|segment| {
            !segment.text.trim().is_empty()
                && (segment.end_secs - timestamp_secs).abs() <= DUPLICATE_TRANSCRIPT_WINDOW_SECS
        })
        .collect();

    nearby.sort_by(|left, right| {
        left.start_secs
            .partial_cmp(&right.start_secs)
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    let mut texts: Vec<String> = Vec::new();
    for start in 0..nearby.len() {
        let mut merged = String::new();
        for end in start..nearby.len().min(start + DUPLICATE_TRANSCRIPT_MERGE_LIMIT) {
            if !merged.is_empty() {
                merged.push(' ');
            }
            merged.push_str(nearby[end].text.trim());
            if !texts.iter().any(|existing| existing == &merged) {
                texts.push(merged.clone());
            }
        }
    }

    texts
}

fn normalize_transcript_text(text: &str) -> String {
    let mut normalized = String::with_capacity(text.len());

    for ch in text.chars() {
        if ch.is_alphanumeric() {
            for lower in ch.to_lowercase() {
                normalized.push(lower);
            }
        } else {
            normalized.push(' ');
        }
    }

    normalized.split_whitespace().collect::<Vec<_>>().join(" ")
}

fn count_common_tokens<'a>(a: &[&'a str], b: &[&'a str]) -> usize {
    let mut counts: HashMap<&'a str, usize> = HashMap::new();
    for token in a {
        *counts.entry(*token).or_insert(0) += 1;
    }

    let mut common = 0usize;
    for token in b {
        if let Some(remaining) = counts.get_mut(*token) {
            if *remaining > 0 {
                *remaining -= 1;
                common += 1;
            }
        }
    }

    common
}

fn longest_common_token_subsequence(a: &[&str], b: &[&str]) -> usize {
    let mut prev = vec![0usize; b.len() + 1];
    let mut curr = vec![0usize; b.len() + 1];

    for token_a in a {
        for (j, token_b) in b.iter().enumerate() {
            curr[j + 1] = if token_a == token_b {
                prev[j] + 1
            } else {
                prev[j + 1].max(curr[j])
            };
        }
        prev.clone_from(&curr);
        curr.fill(0);
    }

    prev[b.len()]
}

fn is_low_signal_token(token: &str) -> bool {
    LOW_SIGNAL_TOKENS
        .iter()
        .any(|low_signal| *low_signal == token)
}

fn meaningful_tokens<'a>(text: &'a str) -> Vec<&'a str> {
    text.split_whitespace()
        .filter(|token| !token.is_empty() && token.len() > 1 && !is_low_signal_token(token))
        .collect()
}

fn transcripts_overlap(a: &str, b: &str) -> bool {
    let normalized_a = normalize_transcript_text(a);
    let normalized_b = normalize_transcript_text(b);
    if normalized_a.is_empty() || normalized_b.is_empty() {
        return false;
    }
    if normalized_a == normalized_b
        || normalized_a.contains(&normalized_b)
        || normalized_b.contains(&normalized_a)
    {
        return true;
    }

    let tokens_a: Vec<&str> = normalized_a.split_whitespace().collect();
    let tokens_b: Vec<&str> = normalized_b.split_whitespace().collect();
    let shorter = tokens_a.len().min(tokens_b.len());
    if shorter < MIN_TOKENS_FOR_OVERLAP {
        return false;
    }

    let common = count_common_tokens(&tokens_a, &tokens_b);
    if common as f64 / shorter as f64 >= TOKEN_COVERAGE_THRESHOLD {
        return true;
    }

    let sequence = longest_common_token_subsequence(&tokens_a, &tokens_b);
    sequence as f64 / shorter as f64 >= TOKEN_SEQUENCE_THRESHOLD
}

fn transcripts_loosely_overlap(a: &str, b: &str) -> bool {
    let normalized_a = normalize_transcript_text(a);
    let normalized_b = normalize_transcript_text(b);
    if normalized_a.is_empty() || normalized_b.is_empty() {
        return false;
    }
    if normalized_a == normalized_b
        || normalized_a.contains(&normalized_b)
        || normalized_b.contains(&normalized_a)
    {
        return true;
    }

    let tokens_a = meaningful_tokens(&normalized_a);
    let tokens_b = meaningful_tokens(&normalized_b);
    let shorter = tokens_a.len().min(tokens_b.len());
    if shorter < MIN_MEANINGFUL_TOKENS_FOR_OVERLAP {
        return false;
    }

    let common = count_common_tokens(&tokens_a, &tokens_b);
    if common as f64 / shorter as f64 >= MEANINGFUL_TOKEN_COVERAGE_THRESHOLD {
        return true;
    }

    let sequence = longest_common_token_subsequence(&tokens_a, &tokens_b);
    sequence as f64 / shorter as f64 >= MEANINGFUL_TOKEN_SEQUENCE_THRESHOLD
}

fn record_speech_window(state: &mut SpeechGateState, rms: f32, peak: f32) {
    state.window_count += 1;
    state.peak_rms = state.peak_rms.max(rms);
    state.peak_amplitude = state.peak_amplitude.max(peak);

    let is_speech_window =
        rms >= SPEECH_WINDOW_RMS_THRESHOLD && peak >= SPEECH_WINDOW_PEAK_THRESHOLD;
    if !is_speech_window {
        state.consecutive_speech_windows = 0;
        return;
    }

    state.speech_window_count += 1;
    state.consecutive_speech_windows += 1;
    state.max_consecutive_speech_windows = state
        .max_consecutive_speech_windows
        .max(state.consecutive_speech_windows);
}

fn speech_gate_decision(state: SpeechGateState, chunk_peak: f32) -> SpeechGateDecision {
    if state.window_count == 0 {
        return SpeechGateDecision {
            skip: false,
            reason: "unavailable",
            peak_rms: state.peak_rms,
            peak_amplitude: state.peak_amplitude,
            window_count: state.window_count,
            speech_window_count: state.speech_window_count,
            max_consecutive_speech_windows: state.max_consecutive_speech_windows,
        };
    }

    let reason = if chunk_peak < FLATLINE_PEAK_THRESHOLD || state.peak_rms < SILENCE_RMS_THRESHOLD {
        Some("silence")
    } else if state.speech_window_count < MIN_SPEECH_FRAMES
        && state.peak_rms < STRONG_SPEECH_RMS_THRESHOLD
        && state.peak_amplitude < STRONG_SPEECH_PEAK_THRESHOLD
    {
        Some("insufficient_speech")
    } else {
        None
    };

    SpeechGateDecision {
        skip: reason.is_some(),
        reason: reason.unwrap_or("speech_detected"),
        peak_rms: state.peak_rms,
        peak_amplitude: state.peak_amplitude,
        window_count: state.window_count,
        speech_window_count: state.speech_window_count,
        max_consecutive_speech_windows: state.max_consecutive_speech_windows,
    }
}

fn evaluate_speech_gate(samples: &[f32]) -> SpeechGateDecision {
    if samples.is_empty() {
        return SpeechGateDecision {
            skip: true,
            reason: "silence",
            peak_rms: 0.0,
            peak_amplitude: 0.0,
            window_count: 0,
            speech_window_count: 0,
            max_consecutive_speech_windows: 0,
        };
    }

    let chunk_peak = samples
        .iter()
        .map(|sample| sample.abs())
        .fold(0.0_f32, f32::max);
    let mut state = SpeechGateState::default();
    for frame in samples.chunks(SPEECH_FRAME_SAMPLES) {
        let len = frame.len().max(1) as f32;
        let rms = (frame.iter().map(|sample| sample * sample).sum::<f32>() / len).sqrt();
        let peak = frame
            .iter()
            .map(|sample| sample.abs())
            .fold(0.0_f32, f32::max);
        record_speech_window(&mut state, rms, peak);
    }

    speech_gate_decision(state, chunk_peak)
}

fn downmix_chunk(samples: Vec<f32>, channels: usize) -> Vec<f32> {
    if channels <= 1 {
        return samples;
    }

    samples
        .chunks(channels)
        .map(|frame| frame.iter().sum::<f32>() / channels as f32)
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    fn segment(start: f64, end: f64, text: &str) -> Segment {
        Segment {
            start,
            end,
            text: text.to_string(),
        }
    }

    #[test]
    fn transcripts_overlap_detects_boundary_repeat() {
        assert!(transcripts_overlap(
            "I need to go to the shops tomorrow",
            "to go to the shops tomorrow"
        ));
    }

    #[test]
    fn loose_overlap_ignores_low_signal_only_match() {
        assert!(!transcripts_loosely_overlap(
            "I think we should do that soon",
            "we should maybe do it soon"
        ));
    }

    #[test]
    fn duplicate_boundary_filter_skips_repeated_opening_segment() {
        let recent_segments = vec![RecentTranscriptSegment {
            start_secs: 10.0,
            end_secs: 12.0,
            text: "I need to go to the shops tomorrow".to_string(),
        }];
        let mut segments = vec![
            segment(0.2, 1.0, "Need to go to the shops tomorrow"),
            segment(1.8, 2.4, "While I am there I need some cheese"),
        ];

        let skipped = filter_duplicate_boundary_segments(&mut segments, 11.8, &recent_segments);

        assert_eq!(skipped, 1);
        assert_eq!(segments.len(), 1);
        assert_eq!(segments[0].text, "While I am there I need some cheese");
    }

    #[test]
    fn remember_recent_segments_prunes_old_history() {
        let mut recent_segments = vec![RecentTranscriptSegment {
            start_secs: 0.0,
            end_secs: 1.0,
            text: "old text".to_string(),
        }];

        remember_recent_segments(
            &mut recent_segments,
            &[segment(0.0, 0.8, "new text")],
            DUPLICATE_HISTORY_RETENTION_SECS + 1.0,
        );

        assert_eq!(recent_segments.len(), 1);
        assert_eq!(recent_segments[0].text, "new text");
    }

    #[test]
    fn speech_gate_treats_near_silence_as_skippable() {
        let samples = vec![0.0004_f32, 0.0002, 0.0003, 0.0001]
            .into_iter()
            .cycle()
            .take(SPEECH_FRAME_SAMPLES * 3)
            .collect::<Vec<_>>();

        let decision = evaluate_speech_gate(&samples);

        assert!(decision.skip);
        assert_eq!(decision.reason, "silence");
    }

    #[test]
    fn speech_gate_rejects_isolated_noise_without_speech_windows() {
        let mut samples = Vec::new();
        for i in 0..(SPEECH_FRAME_SAMPLES * 3) {
            let sample = if i % SPEECH_FRAME_SAMPLES == 0 {
                0.010
            } else {
                0.0011
            };
            samples.push(sample);
        }

        let decision = evaluate_speech_gate(&samples);

        assert!(decision.skip);
        assert_eq!(decision.reason, "insufficient_speech");
        assert_eq!(decision.speech_window_count, 0);
    }

    #[test]
    fn speech_gate_allows_sustained_speech_like_audio() {
        let samples = vec![0.014_f32; SPEECH_FRAME_SAMPLES * 3];

        let decision = evaluate_speech_gate(&samples);

        assert!(!decision.skip);
        assert_eq!(decision.reason, "speech_detected");
        assert_eq!(decision.speech_window_count, 3);
        assert_eq!(decision.max_consecutive_speech_windows, 3);
    }
}