import collections
from math import gcd

import numpy as np
import sounddevice as sd

from .model import SAMPLE_RATE

FRAME_SIZE = 800          # 50ms at 16kHz
PRE_ROLL_FRAMES = 6       # ~0.3s of audio before speech onset
DEFAULT_SILENCE_FRAMES = 16  # ~0.8s of silence to end a segment
SPEECH_ONSET_FRAMES = 3   # ~150ms of speech to trigger
MAX_SPEECH_SECONDS = 30   # force chunk boundary
MIN_LOUD_FRAMES = 8       # need at least ~400ms of loud frames to count as speech


def pause_seconds_to_frames(seconds: float) -> int:
    return max(1, round(seconds / (FRAME_SIZE / SAMPLE_RATE)))


def _query_input(device):
    resolved = device if device is not None else sd.default.device[0]
    info = sd.query_devices(resolved)
    if info["max_input_channels"] < 1:
        raise ValueError(
            f"Device {device!r} ({info['name']}) is not an input device. "
            f"Run `cohere devices` to see current input indices — they can shift between runs on PipeWire."
        )
    return info


def describe_input_device(device) -> str:
    info = _query_input(device)
    return info["name"]


def resolve_input_rate(device) -> int:
    """Pick a samplerate the device will accept. Prefer SAMPLE_RATE; if the device
    refuses (e.g. raw ALSA hw: that doesn't resample), fall back to its native rate."""
    info = _query_input(device)
    try:
        sd.check_input_settings(device=device, samplerate=SAMPLE_RATE, channels=1, dtype="float32")
        return SAMPLE_RATE
    except sd.PortAudioError:
        rate = int(info["default_samplerate"])
        print(f"  Device doesn't support {SAMPLE_RATE}Hz; capturing at {rate}Hz and resampling.")
        return rate


def resample_to_target(audio: np.ndarray, src_rate: int) -> np.ndarray:
    if src_rate == SAMPLE_RATE:
        return audio
    from scipy.signal import resample_poly
    g = gcd(SAMPLE_RATE, src_rate)
    return resample_poly(audio, SAMPLE_RATE // g, src_rate // g).astype(np.float32)


def calibrate_silence(duration=0.5, device=None, compressor=None):
    print("Calibrating silence threshold...")
    rate = resolve_input_rate(device)
    audio = sd.rec(int(duration * rate), samplerate=rate, channels=1, dtype="float32", device=device)
    sd.wait()
    audio = resample_to_target(audio.flatten(), rate)
    if compressor is not None:
        audio = compressor.process(audio)
    rms = np.sqrt(np.mean(audio ** 2))
    threshold = max(rms * 3, 0.01)
    print(f"  Ambient RMS: {rms:.4f}, threshold: {threshold:.4f}")
    return threshold


class VADStateMachine:
    def __init__(self, threshold, silence_frames=DEFAULT_SILENCE_FRAMES):
        self.threshold = threshold
        self.silence_limit = silence_frames
        self.speaking = False
        self.speech_frames = 0
        self.silence_frames = 0
        self.loud_frames = 0
        self.pre_roll = collections.deque(maxlen=PRE_ROLL_FRAMES)
        self.segment = []
        self.segment_start_time = 0.0

    def process_frame(self, frame, elapsed_time):
        """Process one 50ms frame. Returns a (start_time, audio_array) tuple when a
        complete speech segment is detected, otherwise None."""
        rms = np.sqrt(np.mean(frame ** 2))
        is_loud = rms > self.threshold

        if not self.speaking:
            self.pre_roll.append(frame)

            if is_loud:
                self.speech_frames += 1
                if self.speech_frames >= SPEECH_ONSET_FRAMES:
                    self.speaking = True
                    self.silence_frames = 0
                    self.segment = list(self.pre_roll)
                    self.segment_start_time = max(0.0, elapsed_time - len(self.pre_roll) * FRAME_SIZE / SAMPLE_RATE)
                    self.pre_roll = collections.deque(maxlen=PRE_ROLL_FRAMES)
            else:
                self.speech_frames = 0
            return None

        self.segment.append(frame)

        if is_loud:
            self.silence_frames = 0
            self.loud_frames += 1
        else:
            self.silence_frames += 1

        segment_duration = len(self.segment) * FRAME_SIZE / SAMPLE_RATE
        if self.silence_frames >= self.silence_limit or segment_duration >= MAX_SPEECH_SECONDS:
            result = None
            if self.loud_frames >= MIN_LOUD_FRAMES:
                result = (self.segment_start_time, np.concatenate(self.segment))
            self.speaking = False
            self.speech_frames = 0
            self.silence_frames = 0
            self.loud_frames = 0
            self.segment = []
            self.pre_roll = collections.deque(maxlen=PRE_ROLL_FRAMES)
            return result

        return None
