#!/usr/bin/env python3
"""Realtime and file-based ASR processing pipeline using Silero VAD segmentation
and local Orukeet background workers or an OpenAI-compatible ASR server.
"""

import argparse
import copy
import json
import os
import queue
import re
import subprocess
import sys
import tempfile
import threading
import time
import urllib.parse
import urllib.request
import wave
from collections import deque
from pathlib import Path

import numpy as np
import torch
from orukeet import Orukeet

# Thread lock for synchronized access to local engine or shared resources
asr_lock = threading.Lock()


def float32_to_wav_file(float32_array: np.ndarray, sample_rate: int = 16000) -> str:
    """Converts a float32 PCM numpy array into a temporary 16-bit PCM WAV file."""
    int16_pcm = (np.clip(float32_array, -1.0, 1.0) * 32767).astype(np.int16).tobytes()

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_file:
        tmp_filename = tmp_file.name

    with wave.open(tmp_filename, "wb") as wav_out:
        wav_out.setnchannels(1)
        wav_out.setsampwidth(2)
        wav_out.setframerate(sample_rate)
        wav_out.writeframes(int16_pcm)

    return tmp_filename


def transcribe_openai_api(wav_path: str, url: str, model: str) -> dict:
    """Transcribes audio using an OpenAI-compatible ASR server with verbose_json response format."""
    boundary = "----WebKitFormBoundary" + os.urandom(16).hex()

    # Standard OpenAI audio transcription API endpoint target
    endpoint = url
    if not endpoint.endswith("/transcriptions"):
        endpoint = endpoint.rstrip("/") + "/v1/audio/transcriptions"

    with open(wav_path, "rb") as f:
        file_bytes = f.read()

    body = []

    # file field
    body.append(f"--{boundary}".encode("utf-8"))
    body.append(b'Content-Disposition: form-data; name="file"; filename="audio.wav"')
    body.append(b"Content-Type: audio/wav")
    body.append(b"")
    body.append(file_bytes)

    # model field
    body.append(f"--{boundary}".encode("utf-8"))
    body.append(b'Content-Disposition: form-data; name="model"')
    body.append(b"")
    body.append(model.encode("utf-8"))

    # response_format field
    body.append(f"--{boundary}".encode("utf-8"))
    body.append(b'Content-Disposition: form-data; name="response_format"')
    body.append(b"")
    body.append(b"verbose_json")

    body.append(f"--{boundary}--".encode("utf-8"))
    body.append(b"")

    payload = b"\r\n".join(body)

    headers = {
        "Content-Type": f"multipart/form-data; boundary={boundary}",
        "Content-Length": str(len(payload))
    }

    api_key = os.getenv("OPENAI_API_KEY")
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    req = urllib.request.Request(endpoint, data=payload, headers=headers, method="POST")

    with urllib.request.urlopen(req) as response:
        resp_data = response.read().decode("utf-8")
        return json.loads(resp_data)


class SpacyTagger:
    """Optional spaCy integration for NLP tokenization and linguistic tagging."""
    def __init__(self, model_name: str = "fr_core_news_sm"):
        import spacy
        try:
            self.nlp = spacy.load(model_name)
        except Exception as e:
            sys.stderr.write(f"Failed to load spaCy model '{model_name}' ({e}). Attempting fallback download...\n")
            try:
                subprocess.run([sys.executable, "-m", "spacy", "download", model_name], check=True)
                self.nlp = spacy.load(model_name)
            except Exception as err:
                sys.stderr.write(f"Could not load spaCy model '{model_name}': {err}. Disabling spaCy tagging.\n")
                self.nlp = None

    def process_text(self, text: str) -> dict:
        if not self.nlp:
            return {}
        doc = self.nlp(text)
        tokens = [
            {
                "text": token.text,
                "lemma": token.lemma_,
                "pos": token.pos_,
                "tag": token.tag_
            }
            for token in doc
        ]
        return {"tokens": tokens}


class RealtimeVADandOrukeetStreamer:
    """Streams an audio track via ffmpeg/parec, detects boundaries with Silero VAD,
    and enqueues audio turns for Orukeet/OpenAI transcription.
    """
    def __init__(
        self,
        source: str,
        label: str,
        vad_model,
        orukeet_engine: Orukeet,
        min_silence_ms: int,
        track_id: int,
        out_file,
        is_file_input: bool = False,
        spacy_tagger: SpacyTagger = None,
        pre_speech_ms: int = 300,
        api_url: str = None,
        api_model: str = "whisper-1"
    ):
        self.source = source
        self.label = label
        self.vad_model = copy.deepcopy(vad_model)
        self.orukeet_engine = orukeet_engine
        self.api_url = api_url
        self.api_model = api_model
        self.out_file = out_file
        self.track_id = track_id
        self.is_file_input = is_file_input
        self.spacy_tagger = spacy_tagger

        self.sample_rate = 16000
        self.frame_samples = 512  # 32ms frames @ 16kHz
        self.frame_bytes = self.frame_samples * 2

        self.min_silence_samples = int(self.sample_rate * (min_silence_ms / 1000.0))
        self.is_speaking = False
        self.speech_start_sample = 0
        self.speech_start_ts = None
        self.silence_counter_samples = 0

        # Ring buffer for pre-speech audio retention
        self.pre_speech_ms = pre_speech_ms
        pre_speech_frames = max(1, int((pre_speech_ms / 1000.0) * self.sample_rate / self.frame_samples))
        self.pre_speech_ring_buffer = deque(maxlen=pre_speech_frames)

        self.audio_turn_buffer = []

        # Worker queue offloading transcription
        self.asr_queue = queue.Queue()
        self.running = False
        self.read_thread = None
        self.worker_thread = None
        self.process = None

    def emit_event(
        self,
        event_type: str,
        timestamp: float,
        start_ts: int = None,
        end_ts: int = None,
        start: float = None,
        text: str = None,
        phrases: list = None,
        duration: float = None,
        extra: dict = None
    ):
        """Emits structured JSONL event payloads."""
        payload = {
            "label": self.label,
            "event": event_type,
            "timestamp": round(timestamp, 3)
        }
        if start is not None:
            payload["start"] = round(start, 3)
        if start_ts is not None:
            payload["start_ts"] = start_ts
        if end_ts is not None:
            payload["end_ts"] = end_ts
        if text is not None:
            payload["text"] = text
        if phrases:
            payload["phrases"] = phrases
        if duration is not None:
            payload["duration"] = round(duration, 3)
        if extra:
            payload.update(extra)

        self.out_file.write(json.dumps(payload, ensure_ascii=False) + "\n")
        self.out_file.flush()

    def _asr_worker_loop(self):
        """Worker thread processing completed audio turns."""
        while self.running or not self.asr_queue.empty():
            try:
                task = self.asr_queue.get(timeout=0.1)
            except queue.Empty:
                continue

            pcm_data, start_sec, end_sec, duration, start_ts, end_ts = task
            wav_path = float32_to_wav_file(pcm_data, self.sample_rate)

            try:
                if self.api_url:
                    result = transcribe_openai_api(wav_path, self.api_url, self.api_model)
                else:
                    with asr_lock:
                        result = self.orukeet_engine.transcribe(Path(wav_path))

                text = ""
                phrases_info = []

                if isinstance(result, dict):
                    text = result.get("text", "").strip()

                    raw_phrases = result.get("phrases") or result.get("segments") or []
                    for phrase in raw_phrases:
                        p_text = phrase.get("text", "").strip()
                        p_start = round(start_sec + phrase.get("start", 0.0), 3)
                        p_end = round(start_sec + phrase.get("end", 0.0), 3)

                        phrase_dict = {
                            "text": p_text,
                            "start": p_start,
                            "end": p_end
                        }
                        if "words" in phrase:
                            phrase_dict["words"] = phrase["words"]
                        if "tokens" in phrase:
                            phrase_dict["tokens"] = phrase["tokens"]

                        phrases_info.append(phrase_dict)

                elif isinstance(result, str):
                    text = result.strip()

                if text:
                    extra_data = {}
                    if self.spacy_tagger:
                        extra_data["spacy"] = self.spacy_tagger.process_text(text)

                    self.emit_event(
                        "final",
                        timestamp=end_sec,
                        start=start_sec,
                        start_ts=start_ts,
                        end_ts=end_ts,
                        text=text,
                        phrases=phrases_info if phrases_info else None,
                        duration=duration,
                        extra=extra_data
                    )

            except Exception as exc:
                sys.stderr.write(f"[{self.label}] ASR transcription error: {exc}\n")

            finally:
                if os.path.exists(wav_path):
                    os.remove(wav_path)
                self.asr_queue.task_done()

    def _get_subprocess_command(self):
        """Builds reading command depending on input type (file, PulseAudio device, standard ffmpeg)."""
        if self.source.startswith("parec:"):
            device = self.source.replace("parec:", "")
            return ["parec", "--device=" + device, "--format=s16le", "--channels=1", f"--rate={self.sample_rate}"]
        elif self.is_file_input:
            return [
                "ffmpeg",
                "-loglevel", "quiet",
                "-i", self.source,
                "-map", f"0:a:{self.track_id}",
                "-f", "s16le",
                "-ac", "1",
                "-ar", str(self.sample_rate),
                "pipe:1"
            ]
        else:
            return [
                "ffmpeg",
                "-loglevel", "quiet",
                "-i", self.source,
                "-f", "s16le",
                "-ac", "1",
                "-ar", str(self.sample_rate),
                "pipe:1"
            ]

    def _read_loop(self):
        """Streams raw audio samples, evaluates VAD activity, and pushes turns to worker."""
        cmd = self._get_subprocess_command()

        try:
            self.process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
        except Exception as e:
            sys.stderr.write(f"[{self.label}] Failed to start subprocess for '{self.source}': {e}\n")
            return

        if hasattr(self.vad_model, "reset_states"):
            self.vad_model.reset_states()

        total_samples_processed = 0

        while self.running:
            raw_bytes = self.process.stdout.read(self.frame_bytes)
            if len(raw_bytes) < self.frame_bytes:
                break

            audio_int16 = np.frombuffer(raw_bytes, dtype=np.int16)
            float32_frame = audio_int16.astype(np.float32) / 32768.0
            tensor_frame = torch.from_numpy(float32_frame)

            with torch.no_grad():
                speech_prob = self.vad_model(tensor_frame, self.sample_rate).item()

            threshold = 0.5
            current_sec = total_samples_processed / self.sample_rate

            if speech_prob >= threshold:
                self.silence_counter_samples = 0
                if not self.is_speaking:
                    self.is_speaking = True

                    pre_buffer_frames = list(self.pre_speech_ring_buffer)
                    pre_buffer_samples = len(pre_buffer_frames) * self.frame_samples
                    pre_buffer_ms = int((pre_buffer_samples / self.sample_rate) * 1000)

                    self.speech_start_sample = max(0, total_samples_processed - pre_buffer_samples)
                    self.speech_start_ts = int(time.time() * 1000) - pre_buffer_ms

                    self.audio_turn_buffer = []
                    for frame in pre_buffer_frames:
                        self.audio_turn_buffer.extend(frame)

                    actual_start_sec = max(0.0, current_sec - (pre_buffer_samples / self.sample_rate))

                    self.emit_event("start", timestamp=actual_start_sec, start_ts=self.speech_start_ts)

                self.audio_turn_buffer.extend(float32_frame)

            else:
                if self.is_speaking:
                    self.silence_counter_samples += self.frame_samples
                    self.audio_turn_buffer.extend(float32_frame)

                    if self.silence_counter_samples >= self.min_silence_samples:
                        self.is_speaking = False
                        speech_end_ts = int(time.time() * 1000)

                        speech_end_sample = (total_samples_processed + self.frame_samples) - self.silence_counter_samples
                        speech_audio = np.array(self.audio_turn_buffer, dtype=np.float32)

                        start_sec = self.speech_start_sample / self.sample_rate
                        end_sec = max(speech_end_sample / self.sample_rate, start_sec)
                        duration = end_sec - start_sec

                        self.emit_event(
                            "end",
                            timestamp=end_sec,
                            start=start_sec,
                            duration=duration,
                            start_ts=self.speech_start_ts,
                            end_ts=speech_end_ts
                        )

                        self.asr_queue.put((
                            speech_audio,
                            start_sec,
                            end_sec,
                            duration,
                            self.speech_start_ts,
                            speech_end_ts
                        ))

                        self.audio_turn_buffer = []
                        self.pre_speech_ring_buffer.clear()
                        self.silence_counter_samples = 0
                        self.speech_start_ts = None
                else:
                    self.pre_speech_ring_buffer.append(float32_frame)

            total_samples_processed += self.frame_samples

    def start(self):
        self.running = True
        self.read_thread = threading.Thread(target=self._read_loop, daemon=True)
        self.worker_thread = threading.Thread(target=self._asr_worker_loop, daemon=True)
        self.read_thread.start()
        self.worker_thread.start()

    def stop(self):
        self.running = False
        if self.process:
            self.process.terminate()
            self.process.wait()

    def join(self):
        if self.read_thread:
            self.read_thread.join()
        self.asr_queue.join()


def parse_track_args(track_args):
    """Parses track mappings in format `-t 0=speaker1 1=speaker2` or `-t device_name=label` or `-t device_name`."""
    tracks = {}
    if not track_args:
        return {0: "track_0"}

    for item in track_args:
        if "=" in item:
            key, label = item.split("=", 1)
            try:
                tracks[int(key)] = label
            except ValueError:
                tracks[key] = label
        else:
            try:
                tracks[int(item)] = f"track_{item}"
            except ValueError:
                tracks[item] = item
    return tracks


def main():
    default_model = os.getenv("OPENAI_MODEL", "whisper-1")

    parser = argparse.ArgumentParser(description="Realtime VAD and ASR Speech Recognition Pipeline")
    parser.add_argument("-i", "--input", nargs="?", default=None, help="Audio/video file path or stream URL.")
    parser.add_argument("-t", "--track", action="append", nargs="+", help="Audio track or device mapping (e.g. -t dev=label)")
    parser.add_argument("-s", "--silence", "--min-silence-ms", dest="min_silence_ms", type=int, default=500, help="VAD silence threshold in milliseconds")
    parser.add_argument("--pre-speech-ms", type=int, default=300, help="Pre-speech buffer in milliseconds to avoid clipping onset words")
    parser.add_argument("--spacy-model", "--spacy_model", dest="spacy_model", nargs="?", const="fr_core_news_sm", default=None, help="Optional spaCy model (e.g., fr_core_news_sm)")
    parser.add_argument("--installation", type=Path, default=Path("installation.json"), help="Path to Orukeet installation receipt")
    parser.add_argument("--url", type=str, default=None, help="OpenAI-compatible server URL for transcriptions (e.g., http://localhost:8000)")
    parser.add_argument("--model", type=str, default=default_model, help="Model name for OpenAI-compatible server (or legacy Whisper model flag)")
    parser.add_argument("-o", "--output", default=None, help="Output JSONL log path (defaults to stdout)")

    # Legacy CLI compatibility flags
    parser.add_argument("--transcribe", action="store_true", help="Legacy flag (ignored)")
    parser.add_argument("--language", type=str, default=None, help="Legacy language flag (ignored)")

    args = parser.parse_args()

    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8")

    raw_tracks = []
    if args.track:
        for item in args.track:
            if isinstance(item, list):
                raw_tracks.extend(item)
            else:
                raw_tracks.append(item)

    track_mappings = parse_track_args(raw_tracks)
    is_file = os.path.isfile(args.input) if args.input else False

    sys.stderr.write("Loading Silero VAD model...\n")
    vad_model, _ = torch.hub.load(
        repo_or_dir="snakers4/silero-vad",
        model="silero_vad",
        force_reload=False,
        onnx=False
    )

    spacy_tagger = SpacyTagger(args.spacy_model) if args.spacy_model else None
    out_file = open(args.output, "w", encoding="utf-8") if args.output else sys.stdout

    def run_pipeline(asr_engine=None):
        streamers = []
        for track_key, label in track_mappings.items():
            if args.input:
                source = args.input
                track_id = track_key if isinstance(track_key, int) else 0
            else:
                pa_device = str(track_key) if isinstance(track_key, str) else label
                source = f"parec:{pa_device}"
                track_id = 0

            sys.stderr.write(f"[*] Monitoring source: '{source}' (label: '{label}')\n")

            streamer = RealtimeVADandOrukeetStreamer(
                source=source,
                label=label,
                vad_model=vad_model,
                orukeet_engine=asr_engine,
                min_silence_ms=args.min_silence_ms,
                track_id=track_id,
                out_file=out_file,
                is_file_input=is_file,
                spacy_tagger=spacy_tagger,
                pre_speech_ms=args.pre_speech_ms,
                api_url=args.url,
                api_model=args.model
            )
            streamers.append(streamer)

        for streamer in streamers:
            streamer.start()

        for streamer in streamers:
            streamer.join()

    try:
        if args.url:
            sys.stderr.write(f"Using OpenAI-compatible server at {args.url} (model: {args.model})\n")
            run_pipeline(asr_engine=None)
        else:
            config_path = args.installation
            if not config_path.is_file():
                sys.stderr.write(f"Orukeet installation config not found: {config_path}\n")
                sys.exit(1)

            try:
                config = json.loads(config_path.read_text(encoding="utf-8-sig"))
            except Exception as e:
                sys.stderr.write(f"Failed to parse Orukeet installation JSON: {e}\n")
                sys.exit(1)

            sys.stderr.write("Initializing Orukeet engine...\n")
            with Orukeet(config["model"], config["runtime"], device=config["device"]) as asr:
                run_pipeline(asr_engine=asr)

    except Exception as exc:
        sys.stderr.write(f"Execution error: {exc}\n")
        sys.exit(1)
    finally:
        if args.output and out_file != sys.stdout:
            out_file.close()


if __name__ == "__main__":
    main()