#!/usr/bin/env python3
"""
transcriber.py — 単一wavファイルを受け取り文字起こししてjsonlに追記する単一責務モジュール。
録音: recorder.py / ファイル監視: watcher.py が担当。

自己増殖防止：3層dir分離・ファイルパターン除外・whisper上限5・処理済み移動
"""
from __future__ import annotations

import argparse
import fcntl
import glob
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

# ─────────────────────────────────────────────
# 設定
# ─────────────────────────────────────────────
LEFT_NAME  = os.environ.get("LEFT_NAME",  "りつは")
RIGHT_NAME = os.environ.get("RIGHT_NAME", "暦")
THREADS    = int(os.environ.get("THREADS", "4"))

# whisper同時起動上限
WHISPER_MAX_PROCS = 5
WHISPER_WAIT_SEC  = 2      # 待機間隔
WHISPER_TIMEOUT   = 30     # 最大待機秒

# ノイズ語（whisper幻覚フィルタ）
NOISE_PHRASES = [
    "ご視聴ありがとうございました",
    "ありがとうございました",
    "チャンネル登録",
    "高評価",
    "字幕",
    r"^\s*$",       # 空白のみ
    r"^[。、！？…]+$",  # 句読点のみ
]
_NOISE_RE = [re.compile(p) for p in NOISE_PHRASES]

# 入力ファイル名の正規パターン（CHUNK_NNN.wav のみ受理）
VALID_INPUT_RE = re.compile(r"^CHUNK_\d{3,}\.wav$")

# 入力として拒否するパターン
REJECT_PATTERNS = [
    re.compile(r"_L\.wav$"),
    re.compile(r"_R\.wav$"),
    re.compile(r"_processed_"),
]

# ─────────────────────────────────────────────
# パス解決
# ─────────────────────────────────────────────

def resolve_base(input_path: Path) -> Path:
    """recording/raw/CHUNK_*.wav から v2/ ベースを逆引き。"""
    # input_path が .../live_session/v2/recording/raw/CHUNK_NNN.wav の場合
    try:
        idx = input_path.parts.index("v2")
        return Path(*input_path.parts[:idx + 1])
    except ValueError:
        # フォールバック: raw の2階層上
        return input_path.parent.parent.parent


def split_dir(base: Path) -> Path:
    return base / "recording" / "split"


def processed_dir(base: Path) -> Path:
    return base / "recording" / "processed"


def transcripts_dir(base: Path) -> Path:
    return base / "transcripts"


def today_jsonl(base: Path) -> Path:
    date_str = datetime.now().strftime("%Y-%m-%d")
    return transcripts_dir(base) / f"{date_str}.jsonl"

# ─────────────────────────────────────────────
# モデル探索
# ─────────────────────────────────────────────

MODEL_SEARCH_DIRS = [
    Path.home() / ".whisper" / "models",
    Path.home() / "models",
    Path("/usr/local/share/whisper/models"),
    Path("/opt/homebrew/share/whisper/models"),
    # live_session/models も探す
]

MODEL_PRIORITY = ["medium", "small", "base", "tiny"]


def find_model() -> Optional[Path]:
    # 環境変数で明示されていればそれを使う
    env_model = os.environ.get("MODEL_PATH")
    if env_model:
        p = Path(env_model)
        if p.exists():
            return p
        print(f"[transcriber] ERROR: MODEL_PATH={env_model} が見つかりません", file=sys.stderr)
        return None

    # live_session/models も動的に追加
    script_dir = Path(__file__).parent
    extra_dirs = [
        script_dir.parent / "models",   # live_session/models
        script_dir / "models",
    ]
    search_dirs = MODEL_SEARCH_DIRS + extra_dirs

    for priority in MODEL_PRIORITY:
        for d in search_dirs:
            pattern = str(d / f"*{priority}*")
            matches = glob.glob(pattern)
            for m in matches:
                p = Path(m)
                if p.is_file() and p.suffix in (".bin", ".gguf"):
                    return p
    return None

# ─────────────────────────────────────────────
# whisperバイナリ探索
# ─────────────────────────────────────────────

def find_whisper_bin() -> Optional[str]:
    for name in ("whisper-cli", "whisper-cpp", "main"):
        found = shutil.which(name)
        if found:
            return found
    return None

# ─────────────────────────────────────────────
# whisperプロセス上限ガード
# ─────────────────────────────────────────────

def count_whisper_procs() -> int:
    try:
        result = subprocess.run(
            ["pgrep", "-c", "-f", "whisper"],
            capture_output=True, text=True
        )
        count = int(result.stdout.strip()) if result.stdout.strip().isdigit() else 0
        return count
    except Exception:
        return 0


def wait_for_whisper_slot(dry_run: bool = False) -> bool:
    """whisperスロットが空くまで待機。タイムアウトしたらFalseを返す。"""
    if dry_run:
        count = count_whisper_procs()
        print(f"[transcriber] whisperプロセス数: {count} (上限: {WHISPER_MAX_PROCS})")
        return True

    deadline = time.time() + WHISPER_TIMEOUT
    while time.time() < deadline:
        count = count_whisper_procs()
        if count < WHISPER_MAX_PROCS:
            return True
        wait = 1.0 + (hash(time.time()) % 2000) / 1000.0  # 1〜3秒ランダム
        print(f"[transcriber] whisperプロセス数 {count} >= {WHISPER_MAX_PROCS}、{wait:.1f}s待機...", file=sys.stderr)
        time.sleep(wait)
    print(f"[transcriber] SKIP: {WHISPER_TIMEOUT}s待っても空きなし。このチャンクは諦めます。", file=sys.stderr)
    return False

# ─────────────────────────────────────────────
# 入力バリデーション
# ─────────────────────────────────────────────

def validate_input(input_path: Path) -> bool:
    fname = input_path.name
    if not VALID_INPUT_RE.match(fname):
        print(f"[transcriber] REJECT: '{fname}' は CHUNK_NNN.wav パターン外", file=sys.stderr)
        return False
    for pat in REJECT_PATTERNS:
        if pat.search(fname):
            print(f"[transcriber] REJECT: '{fname}' は除外パターンに該当", file=sys.stderr)
            return False
    return True

# ─────────────────────────────────────────────
# ノイズフィルタ
# ─────────────────────────────────────────────

def is_noise(text: str) -> bool:
    for pat in _NOISE_RE:
        if pat.search(text):
            return True
    return False

# ─────────────────────────────────────────────
# ffmpegチャンネル分離
# ─────────────────────────────────────────────

def split_channels(input_wav: Path, split_d: Path) -> tuple[Path, Path]:
    stem = input_wav.stem   # CHUNK_NNN
    left_wav  = split_d / f"{stem}_L.wav"
    right_wav = split_d / f"{stem}_R.wav"

    # L チャンネル
    subprocess.run(
        ["ffmpeg", "-y", "-i", str(input_wav),
         "-filter_complex", "channelsplit=channel_layout=stereo:channels=FL",
         "-c:a", "pcm_s16le", str(left_wav)],
        check=True, capture_output=True
    )
    # R チャンネル
    subprocess.run(
        ["ffmpeg", "-y", "-i", str(input_wav),
         "-filter_complex", "channelsplit=channel_layout=stereo:channels=FR",
         "-c:a", "pcm_s16le", str(right_wav)],
        check=True, capture_output=True
    )
    return left_wav, right_wav

# ─────────────────────────────────────────────
# whisper実行
# ─────────────────────────────────────────────

def run_whisper(bin_path: str, wav_path: Path, model_path: Path) -> str:
    """whisper-cliを実行してテキストを返す。"""
    with tempfile.TemporaryDirectory() as tmpdir:
        out_base = Path(tmpdir) / "out"
        cmd = [
            bin_path,
            "-m", str(model_path),
            "-f", str(wav_path),
            "-t", str(THREADS),
            "--output-txt",
            "--output-file", str(out_base),
            "--no-timestamps",
            "-l", "ja",
        ]
        subprocess.run(cmd, check=True, capture_output=True)
        txt_file = Path(tmpdir) / "out.txt"
        if txt_file.exists():
            return txt_file.read_text(encoding="utf-8").strip()
        return ""

# ─────────────────────────────────────────────
# jsonl排他追記
# ─────────────────────────────────────────────

def append_jsonl(jsonl_path: Path, records: list[dict]) -> None:
    """mkdir lockfileによる排他追記（macOS互換）"""
    lock_path = jsonl_path.parent / f".{jsonl_path.name}.lock"
    deadline = time.time() + 10
    acquired = False
    while time.time() < deadline:
        try:
            lock_path.mkdir(exist_ok=False)
            acquired = True
            break
        except FileExistsError:
            time.sleep(0.1)

    if not acquired:
        print(f"[transcriber] WARN: lockfile取得タイムアウト。強制書き込みします。", file=sys.stderr)

    try:
        with open(jsonl_path, "a", encoding="utf-8") as f:
            fcntl.flock(f, fcntl.LOCK_EX)
            for rec in records:
                f.write(json.dumps(rec, ensure_ascii=False) + "\n")
            fcntl.flock(f, fcntl.LOCK_UN)
    finally:
        if acquired:
            try:
                lock_path.rmdir()
            except Exception:
                pass

# ─────────────────────────────────────────────
# メイン処理
# ─────────────────────────────────────────────

def transcribe(input_path: Path, dry_run: bool = False) -> int:
    """
    Returns:
        0: 成功
        1: エラー
        2: スキップ（上限超過・ノイズのみ等）
    """
    # ── バリデーション ──
    if not validate_input(input_path):
        return 1

    if not input_path.exists():
        print(f"[transcriber] ERROR: ファイルが存在しません: {input_path}", file=sys.stderr)
        return 1

    base   = resolve_base(input_path)
    split_d     = split_dir(base)
    processed_d = processed_dir(base)
    transcripts_d = transcripts_dir(base)

    # ── ディレクトリ確認/作成 ──
    for d in [split_d, processed_d, transcripts_d]:
        d.mkdir(parents=True, exist_ok=True)

    # ── whisperバイナリ探索 ──
    whisper_bin = find_whisper_bin()
    if not whisper_bin:
        print("[transcriber] ERROR: whisper-cli / whisper-cpp が見つかりません", file=sys.stderr)
        return 1

    # ── モデル探索 ──
    model_path = find_model()
    if not model_path:
        print("[transcriber] ERROR: whisperモデルが見つかりません (medium→small→base→tiny 探索失敗)", file=sys.stderr)
        return 1

    if dry_run:
        print(f"[transcriber] dry-run: input={input_path.name} ✓")
        print(f"[transcriber] dry-run: whisper={whisper_bin} ✓")
        print(f"[transcriber] dry-run: model={model_path} ✓")
        print(f"[transcriber] dry-run: split_dir={split_d} ✓")
        print(f"[transcriber] dry-run: processed_dir={processed_d} ✓")
        print(f"[transcriber] dry-run: transcripts_dir={transcripts_d} ✓")
        wait_for_whisper_slot(dry_run=True)
        print("[transcriber] dry-run: 全項目OK")
        return 0

    # ── whisperスロット確認 ──
    if not wait_for_whisper_slot():
        return 2

    # ── ffmpegチャンネル分離 ──
    stem = input_path.stem
    left_wav  = split_d / f"{stem}_L.wav"
    right_wav = split_d / f"{stem}_R.wav"

    try:
        left_wav, right_wav = split_channels(input_path, split_d)
    except subprocess.CalledProcessError as e:
        print(f"[transcriber] ERROR: ffmpegチャンネル分離失敗: {e.stderr.decode()}", file=sys.stderr)
        return 1

    # ── whisper並列実行（L/R最大2並列）──
    records: list[dict] = []
    ts_now = datetime.now(timezone.utc).isoformat()

    procs: list[tuple[subprocess.Popen, str, str]] = []  # (proc, side, speaker)
    for wav_file, speaker in [(left_wav, LEFT_NAME), (right_wav, RIGHT_NAME)]:
        if not wait_for_whisper_slot():
            print(f"[transcriber] SKIP: {speaker} チャンネルはスロット不足でスキップ", file=sys.stderr)
            continue
        with tempfile.TemporaryDirectory() as tmpdir:
            out_base = Path(tmpdir) / "out"
            cmd = [
                whisper_bin,
                "-m", str(model_path),
                "-f", str(wav_file),
                "-t", str(THREADS),
                "--output-txt",
                "--output-file", str(out_base),
                "--no-timestamps",
                "-l", "ja",
            ]
            # 直接実行（並列を安全に扱うため順次）
            try:
                result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
                txt_file = Path(tmpdir) / "out.txt"
                text = txt_file.read_text(encoding="utf-8").strip() if txt_file.exists() else ""
                if text and not is_noise(text):
                    records.append({
                        "ts": ts_now,
                        "speaker": speaker,
                        "text": text,
                        "src_chunk": input_path.name,
                    })
            except subprocess.TimeoutExpired:
                print(f"[transcriber] WARN: {speaker}チャンネルのwhisperがタイムアウト", file=sys.stderr)

    # ── 中間wav削除（split/）──
    for wav_file in [left_wav, right_wav]:
        try:
            wav_file.unlink(missing_ok=True)
        except Exception:
            pass

    # ── jsonl追記 ──
    if records:
        jsonl_path = today_jsonl(base)
        append_jsonl(jsonl_path, records)
        print(f"[transcriber] {len(records)}件追記 → {jsonl_path.name}")
    else:
        print(f"[transcriber] テキストなし（ノイズのみ or 無音）: {input_path.name}")

    # ── 処理済みファイルを移動（再処理防止）──
    # watcher.py が呼び出し前に raw/ → processed/ へ move 済の場合は no-op。
    # 単体実行（watcher 経由でない）時のみ move する。
    dest = processed_d / input_path.name
    try:
        if input_path.resolve() == dest.resolve():
            print(f"[transcriber] 既に processed/ に在り（watcher経由）: {input_path.name}")
        else:
            shutil.move(str(input_path), str(dest))
            print(f"[transcriber] 移動: {input_path.name} → processed/")
    except Exception as e:
        print(f"[transcriber] move スキップ ({e}): {input_path.name}")

    return 0

# ─────────────────────────────────────────────
# dry-run専用モード（入力パスなし）
# ─────────────────────────────────────────────

def dry_run_no_input() -> int:
    """入力ファイルなしのdry-run。環境チェックのみ。"""
    print("[transcriber] dry-run (no input): 環境チェック開始")
    ok = True

    # whisperバイナリ
    whisper_bin = find_whisper_bin()
    if whisper_bin:
        print(f"[transcriber] dry-run: whisper={whisper_bin} ✓")
    else:
        print("[transcriber] dry-run: whisper-cli/whisper-cpp が見つかりません ✗", file=sys.stderr)
        ok = False

    # モデル
    model_path = find_model()
    if model_path:
        print(f"[transcriber] dry-run: model={model_path} ✓")
    else:
        print("[transcriber] dry-run: whisperモデルが見つかりません (medium→small→base→tiny) ✗", file=sys.stderr)
        ok = False

    # ディレクトリ（スクリプトと同階層のv2/を基準）
    base = Path(__file__).parent
    for label, d in [
        ("split_dir",      split_dir(base)),
        ("processed_dir",  processed_dir(base)),
        ("transcripts_dir",transcripts_dir(base)),
    ]:
        if d.exists() or (d.parent.exists()):
            print(f"[transcriber] dry-run: {label}={d} ✓")
        else:
            print(f"[transcriber] dry-run: {label}={d} 作成不可 ✗", file=sys.stderr)
            ok = False

    # whisperプロセス数チェック
    count = count_whisper_procs()
    print(f"[transcriber] dry-run: whisperプロセス数={count} (上限{WHISPER_MAX_PROCS}) ✓")

    # 入力バリデーション関数テスト
    dummy = Path("CHUNK_001.wav")
    assert VALID_INPUT_RE.match(dummy.name), "VALID_INPUT_RE バグ"
    assert not VALID_INPUT_RE.match("CHUNK_001_L.wav"), "VALID_INPUT_RE バグ (L)"
    print("[transcriber] dry-run: VALID_INPUT_RE関数テスト ✓")

    if ok:
        print("[transcriber] dry-run: 全項目OK")
        return 0
    else:
        print("[transcriber] dry-run: 一部項目NG（上記エラー参照）")
        return 1

# ─────────────────────────────────────────────
# エントリポイント
# ─────────────────────────────────────────────

def main() -> None:
    parser = argparse.ArgumentParser(
        description="単一wavを文字起こししてjsonlに追記する"
    )
    parser.add_argument("input_wav", nargs="?", help="recording/raw/CHUNK_NNN.wav のパス")
    parser.add_argument("--dry-run", action="store_true", help="実行せず環境チェックのみ")
    args = parser.parse_args()

    if args.dry_run and not args.input_wav:
        sys.exit(dry_run_no_input())

    if not args.input_wav:
        parser.error("input_wav を指定してください（または --dry-run のみで環境チェック）")

    input_path = Path(args.input_wav).resolve()

    if args.dry_run:
        # ダミーパスでもバリデーション確認
        sys.exit(transcribe(input_path, dry_run=True))

    sys.exit(transcribe(input_path, dry_run=False))


if __name__ == "__main__":
    main()

# ─────────────────────────────────────────────
# 自己増殖防止設計メモ
# ─────────────────────────────────────────────
# 1. 3層dir分離:
#    入力  → recording/raw/CHUNK_NNN.wav    （recorder.pyが生成）
#    中間  → recording/split/CHUNK_NNN_L/R.wav （このスクリプトが生成・処理後即削除）
#    出力  → transcripts/YYYY-MM-DD.jsonl   （このスクリプトが追記）
#    処理済 → recording/processed/CHUNK_NNN.wav （処理後に入力を移動）
#
# 2. ファイルパターン除外:
#    CHUNK_NNN.wav のみ受理。_L.wav / _R.wav / _processed_* は拒否。
#
# 3. whisper上限5:
#    起動前に pgrep -f whisper でカウント。5以上なら最大30秒待機してスキップ。
#
# 4. 処理済み移動:
#    transcribe完了後に raw/ → processed/ に shutil.move。
#    watcher.pyはprocessed/を監視対象外にすること。
