#!/usr/bin/env python3
"""
run_inference.py: Production-ready Korean STT Inference Engine for Notesoop-STT-Ko-v1.

- Google EmbeddingGemma-2 Frozen 300M Audio Tower (MRL 256-dim feature extraction)
- 10-Layer Conformer Reader (21.5M params) on Apple Silicon MPS / CUDA / CPU
- 25,000 Morpheme Vocabulary with Cosine Similarity Matching
- Asynchronous KenLM 3-Gram Language Model Beam Search Rescoring
"""

from __future__ import annotations

import argparse
import difflib
import os
import re
import subprocess
import sys
import time
from pathlib import Path
from typing import Optional, List, Tuple
from concurrent.futures import ThreadPoolExecutor, Future

_ROOT = Path(__file__).resolve().parent
_SRC_DIR = _ROOT / "src"
if str(_SRC_DIR) not in sys.path:
    sys.path.insert(0, str(_SRC_DIR))
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))

import numpy as np
import torch
import yaml

from dataset import KoreanMorphemeTokenizer, normalize_text
from model import EmbeddingGemmaCTCReader
from step2_extract_features import GenuineAudioFeatureExtractor
from decoder import KenLMBeamSearchDecoder


def load_audio(
    file_path: str,
    target_sr: int = 16000,
    start_sec: float = 0.0,
    duration_limit: float = 0.0
) -> Tuple[torch.Tensor, float]:
    """
    Loads any audio container (.wav, .m4a, .mp3, .flac) into 16kHz float32 mono tensor via ffmpeg.
    Uses ultra-fast seeking (-ss / -t) directly in FFmpeg to eliminate reading massive containers into memory.
    """
    p = Path(file_path)
    if not p.exists():
        raise FileNotFoundError(f"Audio file not found: {file_path}")

    cmd = ["ffmpeg", "-nostdin", "-threads", "0"]
    if start_sec > 0:
        cmd.extend(["-ss", str(start_sec)])
    if duration_limit > 0:
        cmd.extend(["-t", str(duration_limit)])
    cmd.extend([
        "-i", str(p),
        "-f", "f32le", "-ac", "1", "-ar", str(target_sr), "-"
    ])
    try:
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
        out, _ = proc.communicate()
        if proc.returncode != 0 or len(out) == 0:
            raise RuntimeError("FFmpeg decode returned empty data")
        arr = np.frombuffer(out, dtype=np.float32).copy()
    except Exception:
        # Fallback to soundfile if ffmpeg is not installed
        import soundfile as sf
        arr, sr = sf.read(str(p), dtype="float32")
        if arr.ndim > 1:
            arr = arr.mean(axis=1)
        if start_sec > 0:
            arr = arr[int(start_sec * sr):]
        if duration_limit > 0:
            arr = arr[:int(duration_limit * sr)]
        if sr != target_sr:
            import math
            gcd = math.gcd(sr, target_sr)
            up = target_sr // gcd
            down = sr // gcd
            import scipy.signal as sps
            arr = sps.resample_poly(arr, up, down).astype(np.float32)

    wav = torch.from_numpy(arr).unsqueeze(0)
    duration_sec = wav.size(-1) / float(target_sr)
    return wav, duration_sec


def merge_chunk_transcripts(chunks: List[str]) -> str:
    """
    Merges overlapping transcript segments by finding maximum character/word overlap.
    Prevents stuttering/duplication across sliding audio window boundaries.
    """
    if not chunks:
        return ""
    if len(chunks) == 1:
        return chunks[0]

    merged = chunks[0]
    for nxt in chunks[1:]:
        if not nxt:
            continue
        max_search = min(len(merged), len(nxt), 40)
        found_overlap = 0

        # Exact substring suffix-prefix match
        for k in range(max_search, 2, -1):
            if merged.endswith(nxt[:k]):
                found_overlap = k
                break

        # Fuzzy sequence matcher
        if found_overlap == 0 and len(merged) >= 10 and len(nxt) >= 10:
            tail = merged[-30:]
            head = nxt[:30]
            matcher = difflib.SequenceMatcher(None, tail, head)
            match = matcher.find_longest_match(0, len(tail), 0, len(head))
            if match.size >= 4 and match.a + match.size == len(tail) and match.b == 0:
                found_overlap = match.size

        if found_overlap > 0:
            merged = merged + nxt[found_overlap:]
        else:
            merged = merged + " " + nxt

    return re.sub(r"\s+", " ", merged).strip()


class EmbeddingGemmaSTT:
    """
    High-performance, standalone Korean STT Pipeline for Notesoop-STT-Ko-v1.
    """
    def __init__(
        self,
        checkpoint_path: Optional[str] = None,
        config_path: Optional[str] = None,
        vocab_path: Optional[str] = None,
        vector_bank_path: Optional[str] = None,
        lm_path: Optional[str] = None,
        device: str = "auto",
        use_kenlm: bool = True,
        lm_alpha: float = 0.12,
        lm_beta: float = 0.40,
        **kwargs  # Ignore legacy kwargs like engine
    ):
        self.root = _ROOT
        self.lm_alpha = lm_alpha
        self.lm_beta = lm_beta

        # 0. Determine Device
        if device == "auto":
            dev_str = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
        else:
            dev_str = "mps" if (device == "metal" and torch.backends.mps.is_available()) else device

        self.device = torch.device(dev_str)

        # 1. Resolve paths & Tokenizer
        v_p = Path(vocab_path) if vocab_path else self.root / "data/vocab/vocab_25k.json"
        if not v_p.exists():
            v_p = self.root / "data/vocab/vocab.json"
        self.tokenizer = KoreanMorphemeTokenizer(str(v_p))

        vb_p = Path(vector_bank_path) if vector_bank_path else self.root / "data/vocab/vector_bank_25k.pt"
        if not vb_p.exists():
            vb_p = self.root / "data/vocab/vector_bank.pt"
        self.vector_bank = torch.load(vb_p, map_location="cpu", weights_only=True)

        # 2. Initialize Conformer Reader
        if checkpoint_path:
            ckpt_p = Path(checkpoint_path)
        else:
            p_inf = self.root / "checkpoints/best_reader_inference.pt"
            p_full = self.root / "checkpoints/best_reader.pt"
            ckpt_p = p_inf if p_inf.exists() else p_full

        if not ckpt_p.exists():
            raise FileNotFoundError(f"Reader checkpoint not found at: {ckpt_p}")

        ckpt = torch.load(ckpt_p, map_location="cpu", weights_only=False)
        state_dict = ckpt.get("model_state_dict", ckpt)

        vb_model = state_dict.get("ctc_head.vector_bank", self.vector_bank)
        d_model = state_dict["reader.input_proj.weight"].shape[0] if "reader.input_proj.weight" in state_dict else 384
        input_dim = state_dict["reader.input_proj.weight"].shape[1] if "reader.input_proj.weight" in state_dict else 256
        detected_layers = len([k for k in state_dict.keys() if k.startswith("reader.layers.") and k.endswith(".final_norm.weight")])
        num_layers = detected_layers if detected_layers > 0 else 10
        dim_ffn = state_dict["reader.layers.0.ffn1.0.weight"].shape[0] if "reader.layers.0.ffn1.0.weight" in state_dict else 1536
        nhead = d_model // 64
        output_dim = state_dict["reader.output_proj.0.weight"].shape[0] if "reader.output_proj.0.weight" in state_dict else 256

        self.reader = EmbeddingGemmaCTCReader(
            input_dim=input_dim,
            output_dim=output_dim,
            d_model=d_model,
            nhead=nhead,
            num_layers=num_layers,
            dim_feedforward=dim_ffn,
            dropout=0.0,
            init_temperature=0.05,
            vector_bank=vb_model,
            use_spec_augment=False,
            trainable_vector_bank=False
        ).to(self.device)
        self.reader.load_state_dict(state_dict, strict=False)
        self.reader.eval()

        # 3. Audio Feature Extractor (Google EmbeddingGemma-2 300M Audio Tower)
        model_name = "google/embeddinggemma-2"
        self.extractor = GenuineAudioFeatureExtractor(device=self.device, model_name=model_name)

        # 4. Optional KenLM Decoder
        self.decoder_kenlm = None
        if use_kenlm:
            lm_file = Path(lm_path) if lm_path else self.root / "data/vocab/korean_lm.bin"
            if lm_file.exists():
                self.decoder_kenlm = KenLMBeamSearchDecoder(
                    tokenizer=self.tokenizer,
                    lm_path=str(lm_file),
                    beam_size=8,
                    alpha=lm_alpha,
                    beta=lm_beta
                )

    @torch.no_grad()
    def transcribe(
        self,
        audio_path: str,
        chunk_sec: float = 16.0,
        overlap_sec: float = 1.0,
        start_sec: float = 0.0,
        duration_limit: float = 0.0,
        batch_size: int = 4
    ) -> Tuple[str, float]:
        """
        Transcribes audio file of arbitrary length into Korean text.
        Optimized with Batched GPU Forward Passes and Asynchronous KenLM Rescoring:
        - CPU Mel Preprocessing
        - GPU 300M Audio Tower Forward (Batched)
        - GPU Conformer Reader + Top-32 Reduction (Batched)
        - Asynchronous KenLM 3-Gram Beam Search Rescoring
        Returns (transcript_text, rtf_speed).
        """
        wav, eval_dur = load_audio(
            audio_path,
            target_sr=16000,
            start_sec=start_sec,
            duration_limit=duration_limit
        )
        sr = 16000
        step_sec = max(chunk_sec - overlap_sec, 0.5) if overlap_sec > 0 else chunk_sec
        step_samples = int(step_sec * sr)
        chunk_samples = int(chunk_sec * sr)
        total_samples = wav.size(-1)

        chunk_ranges: List[Tuple[int, int]] = []
        c_start = 0
        while c_start < total_samples:
            c_end = min(c_start + chunk_samples, total_samples)
            if (c_end - c_start) / sr >= 0.5:
                chunk_ranges.append((c_start, c_end))
            if c_end == total_samples:
                break
            c_start += step_samples

        if not chunk_ranges:
            return "", 0.0

        t0 = time.time()

        # 1. Preprocess Mel-spectrogram on CPU (pad last chunk to chunk_samples for uniform shapes)
        preprocessed_chunks = []
        for cs, ce in chunk_ranges:
            c_wav = wav[:, cs:ce]
            if c_wav.size(-1) < chunk_samples:
                c_wav = torch.nn.functional.pad(c_wav, (0, chunk_samples - c_wav.size(-1)))
            preprocessed_chunks.append(self.extractor.preprocess_audio_cpu(c_wav))

        # 2. Pipelined Batched Inference
        chunk_texts: List[str] = []
        kenlm_futures: List[Future] = []
        use_kenlm_async = (self.decoder_kenlm is not None)

        with ThreadPoolExecutor(max_workers=2) as executor:
            for b_idx in range(0, len(preprocessed_chunks), batch_size):
                b_items = preprocessed_chunks[b_idx : b_idx + batch_size]

                # Step A: Extract Audio Tower Features (GPU)
                if len(b_items) > 1 and len(set(it[3] for it in b_items)) == 1:
                    feats = self.extractor.extract_batch_same_bucket(b_items)
                else:
                    feats = [self.extractor.extract_from_preprocessed(*it) for it in b_items]

                # Step B: Conformer Reader Forward (Batched GPU)
                if len(feats) > 1 and len(set(f.size(0) for f in feats)) == 1:
                    feats_batch = torch.stack(feats, dim=0).to(device=self.device, dtype=torch.float32)
                    logits_batch = self.reader(feats_batch)
                else:
                    logits_batch = [
                        self.reader(f.unsqueeze(0).to(device=self.device, dtype=torch.float32)).squeeze(0)
                        for f in feats
                    ]

                # Step C: Decoding (KenLM Beam Search or Greedy CTC)
                for b in range(len(feats)):
                    logits = logits_batch[b]
                    if use_kenlm_async:
                        topk_lprobs, topk_ids = torch.topk(logits, k=32, dim=-1)
                        blank_lprobs = logits[:, 0]
                        fut = executor.submit(
                            self.decoder_kenlm.decode_topk,
                            topk_ids.detach().cpu().to(torch.int32).numpy(),
                            topk_lprobs.detach().cpu().float().numpy(),
                            blank_lprobs.detach().cpu().float().numpy(),
                            alpha=self.lm_alpha,
                            beta=self.lm_beta
                        )
                        kenlm_futures.append(fut)
                    else:
                        pred_ids = torch.argmax(logits, dim=-1).cpu().tolist()
                        txt = self.tokenizer.decode(pred_ids, remove_blank=True, collapse_repeats=True)
                        chunk_texts.append(txt)

            if use_kenlm_async:
                chunk_texts = [fut.result() for fut in kenlm_futures]

        t_elapsed = time.time() - t0
        rtf = t_elapsed / max(eval_dur, 0.001)

        valid_texts = [t.strip() for t in chunk_texts if t.strip()]
        if overlap_sec > 0:
            full_text = merge_chunk_transcripts(valid_texts)
        else:
            full_text = " ".join(valid_texts)

        return full_text, rtf


def main():
    parser = argparse.ArgumentParser(description="Notesoop-STT-Ko-v1 Standalone Inference")
    parser.add_argument("--audio", type=str, default="sample_audio.wav", help="Path to audio file (.wav, .m4a, .mp3, etc.)")
    parser.add_argument("--checkpoint", type=str, default=None, help="Path to checkpoint (.pt)")
    parser.add_argument("--device", type=str, default="auto", choices=["auto", "mps", "cuda", "cpu"], help="Inference device")
    parser.add_argument("--kenlm", action="store_true", default=True, help="Enable KenLM 3-gram language model rescoring")
    parser.add_argument("--no_kenlm", dest="kenlm", action="store_false", help="Disable KenLM (Greedy decoding only)")
    parser.add_argument("--duration", type=float, default=0.0, help="Duration limit in seconds (0 for full audio)")
    parser.add_argument("--chunk_sec", type=float, default=16.0, help="Chunk length in seconds")
    parser.add_argument("--batch_size", type=int, default=4, help="Batch size for parallel processing (default: 4)")
    parser.add_argument("--output", type=str, default=None, help="Optional output .txt file path")
    args = parser.parse_args()

    print("\n" + "=" * 70)
    print("🌲 Notesoop-STT-Ko-v1: Korean Speech Recognition Engine")
    print("=" * 70)
    print(f"  • Audio File : {args.audio}")
    print(f"  • Device     : {args.device.upper()}")
    print(f"  • Batch Size : {args.batch_size}")
    print(f"  • KenLM LM   : {'Enabled (8-beam rescoring)' if args.kenlm else 'Disabled (Greedy)'}")
    print("-" * 70)

    stt = EmbeddingGemmaSTT(
        checkpoint_path=args.checkpoint,
        device=args.device,
        use_kenlm=args.kenlm
    )

    print("\n🚀 Transcribing...")
    text, rtf = stt.transcribe(
        args.audio,
        chunk_sec=args.chunk_sec,
        duration_limit=args.duration,
        batch_size=args.batch_size
    )

    print("\n" + "=" * 70)
    print("📝 [Transcription Result]:")
    print("=" * 70)
    print(text)
    print("-" * 70)
    print(f"⏱️  Speed / RTF: {rtf:.4f} ({1.0 / max(rtf, 1e-4):.1f}x Real-time on {stt.device.type.upper()})")
    print("=" * 70 + "\n")

    if args.output:
        out_p = Path(args.output)
        out_p.parent.mkdir(parents=True, exist_ok=True)
        out_p.write_text(text, encoding="utf-8")
        print(f"💾 Saved transcription to: {out_p}\n")


if __name__ == "__main__":
    main()