"""Python-side proxy for the Swift ScreenCaptureKit helper.

Spawns `Contents/MacOS/lumo-sysaudio-helper` as a subprocess, talks the
length-prefixed binary protocol documented in swift/SysAudioHelper.swift,
and exposes a small synchronous API:

    proxy = SysAudioProxy()
    proxy.wait_ready(timeout=10)   # blocks until macOS screen-recording
                                   # permission grant + capture is live
    proxy.arm()                    # start emitting chunks
    while running:
        chunk_int16 = proxy.poll_int16(max_samples=16000)  # downsampled
        ...                        # do something (write to wav, transcribe)
    proxy.disarm()
    proxy.shutdown()

The helper emits Float32 mono at the source sample rate (typically 48 kHz).
This proxy linearly downsamples to 16 kHz Int16 — Whisper's expected input
— inline as chunks arrive, so callers only deal with consumption-ready
data. Linear (nearest-neighbour) resampling is fine for speech: introduces
inaudible aliasing at speech frequencies but is 100× faster than a proper
polyphase filter and ScreenCaptureKit's source is already band-limited.
"""

from __future__ import annotations

import os
import struct
import subprocess
import sys
import threading
from pathlib import Path
from typing import Optional

import numpy as np


_TARGET_SAMPLE_RATE = 16000  # Whisper's expected input


def _helper_path() -> Path:
    """Locate the Swift binary either in the running .app bundle (production)
    or under build/swift/ (dev). Raises FileNotFoundError if missing."""
    here = Path(__file__).resolve().parent.parent
    # Frozen: .../Contents/Resources/src/__file__ → grandparent is Contents/
    if getattr(sys, "frozen", False):
        candidate = here.parent / "MacOS" / "lumo-sysaudio-helper"
        if candidate.exists():
            return candidate
    # Dev: project root + build/swift/
    candidate = here / "build" / "swift" / "lumo-sysaudio-helper"
    if candidate.exists():
        return candidate
    raise FileNotFoundError(
        f"lumo-sysaudio-helper binary not found. "
        f"In dev: build it with `swiftc -O swift/SysAudioHelper.swift "
        f"-o build/swift/lumo-sysaudio-helper -framework ScreenCaptureKit "
        f"-framework CoreMedia -framework AVFoundation`. "
        f"In a built .app: scripts/build_app.sh should have compiled and "
        f"copied it to Contents/MacOS/."
    )


class SysAudioProxy:
    """One subprocess per instance. Idempotent shutdown.

    `token` (optional): a JWS-format meeting capture token issued by the
    license server. When provided, it's sent to the helper as the first
    command (frame type 'T') before any ARM. The helper verifies it
    against its hardcoded Ed25519 public key and refuses to start
    capture if invalid/expired/wrong-device.
    """

    def __init__(self, token: Optional[str] = None) -> None:
        argv = [str(_helper_path())]
        helper_log = os.path.expanduser("~/Library/Logs/Lirelle_sysaudio_helper.log")
        try:
            os.makedirs(os.path.dirname(helper_log), exist_ok=True)
            stderr_fh = open(helper_log, "a", buffering=1)
        except Exception:
            stderr_fh = subprocess.DEVNULL

        self._proc = subprocess.Popen(
            argv,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=stderr_fh,
            bufsize=0,
        )

        # Audio plumbing: chunks accumulate in _buffer (Int16, 16 kHz) and
        # callers drain via poll_int16(). _ready_event flips when the helper
        # confirms capture is actually live (after any screen-recording
        # permission grant). _error holds the most recent 'E' frame payload
        # so callers can surface a meaningful failure reason in the UI.
        self._buffer = bytearray()
        self._buffer_lock = threading.Lock()
        self._ready_event = threading.Event()
        self._error: Optional[str] = None
        self._stopped = False

        # Source format becomes known after the first chunk. We assume it
        # stays stable (it does — SCK doesn't switch sample rate mid-stream).
        self._src_sample_rate: Optional[int] = None
        self._downsample_residual = 0.0  # fractional-sample carry between chunks

        self._reader_thread = threading.Thread(
            target=self._reader_loop, daemon=True
        )
        self._reader_thread.start()

        # Hand the helper its meeting capture token. Frame format:
        #   [1 byte 'T'][4 LE u32 length][N bytes utf-8 token]
        # Helper verifies Ed25519 signature + exp + machine_id before
        # accepting any subsequent ARM. If we have no token (dev / Free
        # tier), the helper will refuse to start once an ARM arrives —
        # we still spawn it so the same wait_ready() error path applies.
        if token:
            try:
                token_bytes = token.encode("utf-8")
                self._proc.stdin.write(b"T")
                self._proc.stdin.write(struct.pack("<I", len(token_bytes)))
                self._proc.stdin.write(token_bytes)
                self._proc.stdin.flush()
            except (BrokenPipeError, OSError) as e:
                print(f"[sysaudio_proxy] failed to send token: {e}", flush=True)

    # ── subprocess protocol ────────────────────────────────────────────

    def _read_exact(self, n: int) -> Optional[bytes]:
        out = self._proc.stdout
        if out is None:
            return None
        buf = bytearray()
        while len(buf) < n:
            chunk = out.read(n - len(buf))
            if not chunk:
                return None
            buf += chunk
        return bytes(buf)

    def _reader_loop(self) -> None:
        while True:
            hdr = self._read_exact(5)
            if hdr is None:
                break
            ftype = hdr[0:1]
            (plen,) = struct.unpack("<I", hdr[1:5])
            if plen > 0:
                payload = self._read_exact(plen)
                if payload is None:
                    break
            else:
                payload = b""

            if ftype == b"R":
                self._ready_event.set()
            elif ftype == b"C":
                self._handle_chunk(payload)
            elif ftype == b"E":
                msg = payload.decode("utf-8", errors="replace")
                print(f"[sysaudio_proxy] helper error: {msg}", flush=True)
                # First error wins — callers surface it as the failure reason.
                if self._error is None:
                    self._error = msg
                # Unblock wait_ready so the caller doesn't sit on the 30 s
                # timeout when we know up-front that capture won't start
                # (TCC denial, no displays, etc.).
                self._ready_event.set()
            # 'O' (pong) intentionally ignored — we don't currently watchdog
            # this helper (the screen-recording API itself doesn't wedge the
            # way mic CoreAudio does).

    def _send(self, cmd: bytes) -> bool:
        if self._proc is None or self._proc.stdin is None:
            return False
        try:
            self._proc.stdin.write(cmd)
            self._proc.stdin.flush()
            return True
        except (BrokenPipeError, OSError):
            return False

    # ── chunk handling + inline downsampling ───────────────────────────

    def _handle_chunk(self, payload: bytes) -> None:
        """Decode a 'C' frame and downsample to 16 kHz Int16."""
        if len(payload) < 8:
            return
        src_sr = struct.unpack("<I", payload[:4])[0]
        # channel count, bits per channel, is_float at payload[4:7]
        # We hard-assume mono Float32 because that's what the Swift side
        # ships with our current SCStreamConfiguration.channelCount=1.
        audio_bytes = payload[8:]
        if not audio_bytes:
            return

        if self._src_sample_rate is None:
            self._src_sample_rate = src_sr

        samples_f32 = np.frombuffer(audio_bytes, dtype=np.float32)
        if samples_f32.size == 0:
            return

        # Downsample to 16 kHz with carried fractional indexing — preserves
        # alignment across chunks so total duration matches expectation.
        if src_sr == _TARGET_SAMPLE_RATE:
            downsampled = samples_f32
        else:
            ratio = src_sr / _TARGET_SAMPLE_RATE
            # Continuing index from where we left off; subtract any whole-
            # sample advance the residual covered.
            start = self._downsample_residual
            n_out = int((samples_f32.size - start) / ratio) + 1
            if n_out < 1:
                # Chunk too small to produce even one output sample —
                # carry it over. (Rare; SCK chunks are ~20 ms.)
                self._downsample_residual = start - samples_f32.size
                return
            indices = np.arange(n_out) * ratio + start
            mask = indices < samples_f32.size
            indices = indices[mask].astype(np.int64)
            downsampled = samples_f32[indices]
            # Carry fractional position so we don't drift over many chunks.
            next_pos = (n_out * ratio + start) - samples_f32.size
            self._downsample_residual = next_pos

        # Float32 [-1.0, 1.0] → Int16
        clipped = np.clip(downsampled, -1.0, 1.0)
        int16 = (clipped * 32767.0).astype(np.int16)
        with self._buffer_lock:
            self._buffer.extend(int16.tobytes())

    # ── public API ─────────────────────────────────────────────────────

    def wait_ready(self, timeout: float = 5.0) -> bool:
        """Block until the helper signals READY (capture is live) OR errors
        out. Returns True only on a clean ready. On False, read `.error`
        for the failure reason (TCC denial, no displays, etc.). Shorter
        default timeout than before: 5 s is plenty for ScreenCaptureKit's
        startup; longer just dragged a wedged TCC prompt's UI freeze out."""
        self._ready_event.wait(timeout)
        return self._ready_event.is_set() and self._error is None

    @property
    def error(self) -> Optional[str]:
        return self._error

    def arm(self) -> None:
        self._send(b"A")

    def disarm(self) -> None:
        self._send(b"D")

    def poll_int16(self, max_samples: int = 16000) -> bytes:
        """Drain up to max_samples × 2 bytes of accumulated audio. Returns
        empty bytes if no data buffered."""
        max_bytes = max_samples * 2
        with self._buffer_lock:
            if not self._buffer:
                return b""
            take = self._buffer[:max_bytes]
            del self._buffer[:max_bytes]
            return bytes(take)

    @property
    def source_sample_rate(self) -> Optional[int]:
        return self._src_sample_rate

    def shutdown(self) -> None:
        if self._stopped:
            return
        self._stopped = True
        self._send(b"Q")
        if self._proc is not None:
            try:
                self._proc.wait(timeout=2.0)
            except subprocess.TimeoutExpired:
                try:
                    self._proc.kill()
                except Exception:
                    pass
