#!/usr/bin/env python3
"""
FreePBX Hívásátíró – parancssori eszköz
=======================================
Stereo MP3/WAV hívásrögzítések átírása magyar nyelvű, időbélyeges TXT dialógussá.

Telepítés:
    pip install mlx-whisper pydub soundfile numpy pyyaml resemblyzer

Használat:
    python transcribe.py /mappa/hangfajlok
    python transcribe.py /mappa/hangfajlok -c config.yaml
    python transcribe.py /mappa/hangfajlok -r -v
    python transcribe.py /mappa/hangfajlok -m mlx-community/whisper-large-v3-mlx
"""

import os
import sys
import argparse
import logging
from pathlib import Path

import numpy as np
import yaml
from pydub import AudioSegment

# ─── Whisper backend: mlx (Apple Silicon) vagy faster-whisper (fallback) ───
try:
    import mlx_whisper
    from huggingface_hub import snapshot_download
    from huggingface_hub.utils import LocalEntryNotFoundError
    WHISPER_BACKEND = "mlx"
except ImportError:
    try:
        from faster_whisper import WhisperModel
        WHISPER_BACKEND = "faster"
    except ImportError:
        print(
            "Hiányzó csomag! Telepítsd az egyiket:\n"
            "  pip install mlx-whisper          (Apple Silicon – ajánlott)\n"
            "  pip install faster-whisper        (más platform)\n",
            file=sys.stderr,
        )
        sys.exit(1)

# ─── Hangminta-alapú azonosítás (opcionális) ───────────────────────────────
try:
    from resemblyzer import VoiceEncoder, preprocess_wav as _preprocess_wav
    RESEMBLYZER_OK = True
except ImportError:
    RESEMBLYZER_OK = False

TARGET_SR = 16_000  # Whisper natív mintavételi frekvencia


# ═══════════════════════════════════════════════════════════════════════════
# Segédeszközök
# ═══════════════════════════════════════════════════════════════════════════

def setup_logging(verbose: bool) -> None:
    logging.basicConfig(
        level=logging.DEBUG if verbose else logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
        datefmt="%H:%M:%S",
    )


def load_config(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def load_vocabulary(path: str) -> list:
    if not path or not os.path.exists(path):
        return []
    words = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            word = line.strip()
            if word and not word.startswith("#"):
                words.append(word)
    return words


def format_ts(seconds: float) -> str:
    """Másodpercek → MM:SS.mmm  (vagy HH:MM:SS ha > 1 óra)"""
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = int(seconds % 60)
    ms = int((seconds % 1) * 1000)
    if h:
        return f"{h:02d}:{m:02d}:{s:02d}"
    return f"{m:02d}:{s:02d}.{ms:03d}"


# ═══════════════════════════════════════════════════════════════════════════
# Hangfájl betöltés és csatornaszétválasztás
# ═══════════════════════════════════════════════════════════════════════════

def _pydub_to_float32(segment: AudioSegment) -> np.ndarray:
    max_val = float(2 ** (8 * segment.sample_width - 1))
    return np.array(segment.get_array_of_samples(), dtype=np.float32) / max_val


def load_stereo_channels(audio_path: Path) -> tuple:
    """
    Sztereó hangfájl betöltése.  Visszatér: (left, right) – float32 @ 16 kHz.
    JointStereo MP3 esetén az L/R csatornák szét lesznek választva.
    """
    audio = AudioSegment.from_file(str(audio_path))
    audio = audio.set_frame_rate(TARGET_SR)  # resample → 16 kHz

    if audio.channels == 1:
        logging.warning("Mono fájl – mindkét csatorna azonos lesz")
        data = _pydub_to_float32(audio)
        return data, data.copy()

    # Sztereó → bal + jobb mono
    ch = audio.split_to_mono()
    return _pydub_to_float32(ch[0]), _pydub_to_float32(ch[1])


def load_voice_sample(path: str) -> np.ndarray:
    """Hangminta betöltése float32 @ 16 kHz mono tömbként."""
    audio = AudioSegment.from_file(path).set_frame_rate(TARGET_SR).set_channels(1)
    return _pydub_to_float32(audio)


# ═══════════════════════════════════════════════════════════════════════════
# Hangminta-alapú ügyfélszolgálatos-azonosítás
# ═══════════════════════════════════════════════════════════════════════════

def identify_agent_channel(
    left: np.ndarray,
    right: np.ndarray,
    agent_samples: list,   # [(név, np.ndarray), ...]
    encoder: "VoiceEncoder",
    threshold: float = 0.75,
) -> tuple:
    """
    Összehasonlítja a két csatorna hangját a hangmintákkal koszinusz-hasonlóság alapján.
    Visszatér: (ügyfélszolgálatos_neve | None, csatorna_index)
      - Ha egyik csatorna sem éri el a küszöbértéket, a név None lesz (ismeretlen hívás).
      - Küszöbérték (threshold): 0.0–1.0, ajánlott 0.72–0.80.
        Túl alacsony → téves egyezések. Túl magas → valódi egyezéseket hagy ki.
    """
    def embed(audio: np.ndarray) -> np.ndarray:
        wav = _preprocess_wav(audio, source_sr=TARGET_SR)
        return encoder.embed_utterance(wav)

    left_emb  = embed(left)
    right_emb = embed(right)

    best_name  = None
    best_ch    = 0
    best_score = -2.0

    for name, sample_audio in agent_samples:
        sample_emb = embed(sample_audio)
        sim_l = float(np.dot(left_emb,  sample_emb))
        sim_r = float(np.dot(right_emb, sample_emb))
        # Mindig INFO szinten logolunk, hogy a threshold finomhangolásához látszanak az értékek
        logging.info(f"  Hasonlóság – {name}: bal={sim_l:.3f}  jobb={sim_r:.3f}  (küszöb={threshold:.2f})")

        if sim_l > best_score:
            best_score, best_name, best_ch = sim_l, name, 0
        if sim_r > best_score:
            best_score, best_name, best_ch = sim_r, name, 1

    # Csak akkor fogadjuk el az egyezést, ha meghaladja a küszöbértéket
    if best_score < threshold:
        logging.info(
            f"  Egyezés elutasítva (score={best_score:.3f} < küszöb={threshold:.2f}) "
            f"→ alapértelmezett csatorna-beállítás lép életbe"
        )
        return None, best_ch

    side = "bal" if best_ch == 0 else "jobb"
    logging.info(f"  Azonosítva: {best_name} → {side} csatorna  (score={best_score:.3f} ✓)")
    return best_name, best_ch


# ═══════════════════════════════════════════════════════════════════════════
# Whisper átírás
# ═══════════════════════════════════════════════════════════════════════════

def build_initial_prompt(vocabulary: list) -> str:
    base = (
        "Ügyfélszolgálati telefonbeszélgetés. "
        "Internet, kábeltévé, e-mail és domain szolgáltatások. "
        "Magyarországi internetszolgáltató."
    )
    if not vocabulary:
        return base
    # ~150 token határ – az első 60 szó belefér általában
    vocab_str = ", ".join(vocabulary[:60])
    return f"{base} Szakkifejezések: {vocab_str}."


def _transcribe_mlx(audio: np.ndarray, model_name: str, prompt: str) -> list:
    result = mlx_whisper.transcribe(
        audio,
        path_or_hf_repo=model_name,
        language="hu",
        initial_prompt=prompt,
        word_timestamps=False,
        condition_on_previous_text=True,
        temperature=0.0,
        no_speech_threshold=0.6,
    )
    return [
        {
            "start": seg["start"],
            "end":   seg["end"],
            "text":  seg["text"].strip(),
        }
        for seg in result.get("segments", [])
        if seg["text"].strip()
    ]


def _transcribe_faster(audio: np.ndarray, model: "WhisperModel", prompt: str) -> list:
    gen, _ = model.transcribe(
        audio,
        language="hu",
        initial_prompt=prompt,
        condition_on_previous_text=True,
        vad_filter=True,
        vad_parameters={"min_silence_duration_ms": 400},
        temperature=0.0,
    )
    return [
        {
            "start": seg.start,
            "end":   seg.end,
            "text":  seg.text.strip(),
        }
        for seg in gen
        if seg.text.strip()
    ]


def transcribe_channel(
    audio: np.ndarray,
    model_or_name,
    prompt: str,
    speaker_name: str,
) -> list:
    """Egy mono csatorna átírása. Visszatér: [{start, end, text, speaker}, ...]"""
    if WHISPER_BACKEND == "mlx":
        raw = _transcribe_mlx(audio, model_or_name, prompt)
    else:
        raw = _transcribe_faster(audio, model_or_name, prompt)

    for seg in raw:
        seg["speaker"] = speaker_name
    return raw


# ═══════════════════════════════════════════════════════════════════════════
# Szegmensek összefésülése és TXT kimenet
# ═══════════════════════════════════════════════════════════════════════════

def merge_segments(agent_segs: list, customer_segs: list, gap_sec: float = 0.8) -> list:
    """
    Időrend szerint összefésüli a két csatorna szegmenseit.
    Az azonos beszélő egymást követő, közeli (< gap_sec) szegmenseit összevonja.
    """
    all_segs = sorted(agent_segs + customer_segs, key=lambda x: x["start"])
    merged = []
    for seg in all_segs:
        prev = merged[-1] if merged else None
        if (
            prev
            and prev["speaker"] == seg["speaker"]
            and seg["start"] - prev["end"] <= gap_sec
        ):
            prev["end"]   = seg["end"]
            prev["text"] += " " + seg["text"]
        else:
            merged.append(dict(seg))
    return merged


def write_transcript(segments: list, output_path: Path, source_name: str) -> None:
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(f"Forrás: {source_name}\n")
        f.write("=" * 60 + "\n\n")
        for seg in segments:
            f.write(f"[{format_ts(seg['start'])}] {seg['speaker']}: {seg['text']}\n")


# ═══════════════════════════════════════════════════════════════════════════
# Egy hangfájl teljes feldolgozása
# ═══════════════════════════════════════════════════════════════════════════

def process_file(
    audio_path: Path,
    config: dict,
    vocabulary: list,
    model_or_name,
    encoder,
    agent_samples: list,
    suffix: str,
) -> None:
    output_path = audio_path.parent / (audio_path.stem + suffix)

    if output_path.exists():
        logging.info(f"Kész (kihagyva): {audio_path.name}")
        return

    logging.info(f"Feldolgozás: {audio_path.name}")

    # 1. Csatornák szétválasztása
    try:
        left, right = load_stereo_channels(audio_path)
    except Exception as e:
        logging.error(f"  Betöltési hiba: {e}")
        return

    # 2. Ügyfélszolgálatos csatornájának azonosítása
    default_agent   = config.get("default_agent_name",    "Ügyfélszolgálatos")
    default_ch      = config.get("default_agent_channel", 1)   # FreePBX-ben általában jobb = fogadó
    customer_label  = config.get("customer_label",        "Ügyfél")

    threshold = config.get("speaker_similarity_threshold", 0.75)

    if encoder and agent_samples:
        try:
            matched_name, agent_ch = identify_agent_channel(
                left, right, agent_samples, encoder, threshold=threshold
            )
        except Exception as e:
            logging.warning(f"  Azonosítási hiba, alapértelmezett beállítás: {e}")
            matched_name, agent_ch = None, default_ch

        if matched_name is not None:
            agent_name = matched_name
        else:
            # Nincs elég magabiztos egyezés – alapértelmezett névvel és csatornával dolgozunk
            agent_name = default_agent
            agent_ch   = default_ch
            logging.info(f"  Ismeretlen hang → alapértelmezett: {agent_name} / {'bal' if agent_ch == 0 else 'jobb'} csatorna")
    else:
        agent_name, agent_ch = default_agent, default_ch
        side = "bal" if agent_ch == 0 else "jobb"
        logging.info(f"  Alapértelmezett: {agent_name} → {side} csatorna")

    agent_audio    = left  if agent_ch == 0 else right
    customer_audio = right if agent_ch == 0 else left

    # 3. Átírás
    prompt = build_initial_prompt(vocabulary)
    try:
        logging.info(f"  Átírás: {agent_name}...")
        agent_segs = transcribe_channel(agent_audio, model_or_name, prompt, agent_name)

        logging.info(f"  Átírás: {customer_label}...")
        cust_segs  = transcribe_channel(customer_audio, model_or_name, prompt, customer_label)
    except Exception as e:
        logging.error(f"  Átírási hiba: {e}")
        return

    # 4. Összefésülés és mentés
    merged = merge_segments(agent_segs, cust_segs)
    write_transcript(merged, output_path, audio_path.name)
    logging.info(f"  Mentve: {output_path.name}  ({len(merged)} fordulat)")


# ═══════════════════════════════════════════════════════════════════════════
# Modell cache kezelés
# ═══════════════════════════════════════════════════════════════════════════

def _ensure_model_local(model_name: str) -> str:
    """
    Visszaadja a modell helyi cache elérési útját.
    - Ha már le van töltve: azonnal, hálózati kérés nélkül.
    - Ha még nincs: letölti egyszer, utána soha többé nem nyúl a hálózathoz.
    """
    # 1. Próbáljuk offline – ha megvan, egyetlen HTTP kérés sem megy ki
    try:
        path = snapshot_download(repo_id=model_name, local_files_only=True)
        logging.info(f"Modell betöltve (helyi cache): {path}")
        return path
    except LocalEntryNotFoundError:
        pass

    # 2. Nincs még letöltve – letöltjük most egyszer
    logging.info(f"Modell első letöltése: {model_name}  (ezután már csak a cache-ből tölt be)")
    path = snapshot_download(repo_id=model_name, local_files_only=False)
    logging.info(f"Letöltés kész, cache: {path}")
    return path


# ═══════════════════════════════════════════════════════════════════════════
# Main
# ═══════════════════════════════════════════════════════════════════════════

def main() -> None:
    parser = argparse.ArgumentParser(
        description="FreePBX hívásrögzítések átírása magyar TXT dialógussá",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Példák:
  python transcribe.py /Volumes/felvetelek/2024
  python transcribe.py /Volumes/felvetelek/2024 -r               # almappák is
  python transcribe.py /Volumes/felvetelek/2024 -c config.yaml -v
  python transcribe.py /Volumes/felvetelek/2024 -m mlx-community/whisper-large-v3-mlx
        """,
    )
    parser.add_argument("folder",
        help="Hangfájlokat tartalmazó mappa")
    parser.add_argument("--config", "-c", default="config.yaml",
        help="Konfigurációs fájl (alap: config.yaml)")
    parser.add_argument("--model", "-m", default=None,
        help="Whisper modell neve (felülírja a config beállítását)")
    parser.add_argument("--recursive", "-r", action="store_true",
        help="Almappák feldolgozása is")
    parser.add_argument("--suffix", default=None,
        help="Kimeneti fájlnév végzete (alap: _transcript.txt)")
    parser.add_argument("--verbose", "-v", action="store_true",
        help="Részletes log kimenet")
    args = parser.parse_args()

    setup_logging(args.verbose)

    # Konfig betöltése
    if not os.path.exists(args.config):
        logging.error(f"Konfigurációs fájl nem található: {args.config}")
        sys.exit(1)
    config = load_config(args.config)

    # Szótár
    vocab_path = config.get("vocabulary_file", "vocabulary.txt")
    vocabulary = load_vocabulary(vocab_path)
    if vocabulary:
        logging.info(f"Szótár betöltve: {len(vocabulary)} szó")
    else:
        logging.info("Szótár: nem betöltve (elhagyható)")

    # Whisper modell
    model_name = args.model or config.get("whisper_model", "mlx-community/whisper-large-v3-mlx")

    if WHISPER_BACKEND == "mlx":
        model_or_name = _ensure_model_local(model_name)
    else:
        logging.info(f"Whisper backend: faster-whisper  |  modell: {model_name}")
        model_or_name = WhisperModel(model_name, device="cpu", compute_type="int8")

    # Hangminta-azonosító
    encoder = None
    agent_samples = []

    if RESEMBLYZER_OK:
        agents = config.get("agents", [])
        loadable = [a for a in agents if a.get("voice_sample") and os.path.exists(a["voice_sample"])]
        if loadable:
            logging.info("Hangminta-azonosító inicializálása (resemblyzer)...")
            encoder = VoiceEncoder()
            for agent in loadable:
                try:
                    sample = load_voice_sample(agent["voice_sample"])
                    agent_samples.append((agent["name"], sample))
                    logging.info(f"  OK: {agent['name']}")
                except Exception as e:
                    logging.warning(f"  Hiba ({agent['name']}): {e}")
        else:
            logging.info("Nincs hangminta konfigurálva – automatikus azonosítás kikapcsolva")
    else:
        logging.warning(
            "resemblyzer nincs telepítve – hangminta-alapú azonosítás kikapcsolva.\n"
            "  Telepítés: pip install resemblyzer"
        )

    # Hangfájlok összegyűjtése
    folder = Path(args.folder)
    if not folder.exists():
        logging.error(f"Mappa nem létezik: {folder}")
        sys.exit(1)

    audio_files = []
    for pat in ("*.mp3", "*.wav", "*.MP3", "*.WAV", "*.m4a", "*.M4A"):
        fn = folder.rglob if args.recursive else folder.glob
        audio_files.extend(fn(pat))
    audio_files = sorted(set(audio_files))

    if not audio_files:
        logging.warning("Nem találhatók hangfájlok a megadott mappában")
        sys.exit(0)

    suffix = args.suffix or config.get("output_suffix", "_transcript.txt")
    logging.info(f"{len(audio_files)} hangfájl megtalálva")

    # Feldolgozás
    ok = 0
    skip = 0
    err = 0
    for i, path in enumerate(audio_files, 1):
        output_path = path.parent / (path.stem + suffix)
        if output_path.exists():
            skip += 1
        logging.info(f"[{i}/{len(audio_files)}]")
        before_err = err
        try:
            process_file(path, config, vocabulary, model_or_name, encoder, agent_samples, suffix)
            if not output_path.exists():
                err += 1
            elif output_path.stat().st_mtime > path.stat().st_mtime - 5:
                ok += 1
        except Exception as e:
            logging.error(f"Váratlan hiba: {path.name}: {e}")
            err += 1

    logging.info(f"\nFeldolgozás kész.  Feldolgozva: {ok}  Kihagyva: {skip}  Hiba: {err}")


if __name__ == "__main__":
    main()