"""Audio-to-JSONL processor for get笔 recordings.

Runs local ASR + speaker diarization and emits Claude-compatible JSONL.

Usage:
  audio_processor /path/to/recording.m4a

Outputs to stdout one JSONL line per speaker turn, e.g.:
{
  "type": "user",
  "message": {"content": [{"type": "text", "text": "Speaker A: Hello."}]},
  "timestamp": "2026-03-20T14:22:05.123Z",
  "sessionId": "audio_20260320_142205"
}

Then optionally emits:
{
  "type": "assistant",
  "message": {
    "content": [{
      "type": "tool_use",
      "id": "t1",
      "name": "SummarizeConversation",
      "input": { ... }
    }]
  },
  "timestamp": "...",
  "sessionId": "..."
}

Requires:
- faster-whisper (pip install faster-whisper)
- pyannote.audio (pip install pyannote.audio)
- pyannote.database (for demo, optional)
- torch, torchaudio

Model notes:
- Whisper model: "large-v3-turbo" (best speed/quality for Chinese)
- Diarization model: "pyannote/speaker-diarization-3.1" (SOTA, multilingual)
"""

import argparse
import json
import os
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import List, Dict, Any, Optional

import torch
import whisper
from pyannote.audio import Pipeline
from pyannote.core import Segment

# --- Configurable defaults ---
WHISPER_MODEL = "large-v3-turbo"
DIARIZATION_MODEL = "pyannote/speaker-diarization-3.1"
WHISPER_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def load_models():
    """Load ASR and diarization models once."""
    print("Loading ASR model...", file=sys.stderr)
    asr_model = whisper.load_model(WHISPER_MODEL, device=WHISPER_DEVICE)

    print("Loading diarization pipeline...", file=sys.stderr)
    # Requires HF_TOKEN in env for pyannote (free API key at huggingface.co/settings/tokens)
    if not os.getenv("HF_TOKEN"):
        raise RuntimeError("HF_TOKEN environment variable required for pyannote.audio")
    diarization_pipeline = Pipeline.from_pretrained(DIARIZATION_MODEL)
    diarization_pipeline.to(torch.device(WHISPER_DEVICE))

    return asr_model, diarization_pipeline


def transcribe_and_diarize(
    audio_path: str,
    asr_model,
    diarization_pipeline,
    min_speaker_duration: float = 0.5,
) -> List[Dict[str, Any]]:
    """Run ASR + diarization and return list of {'speaker', 'text', 'start', 'end'}"""
    audio_path = str(audio_path)

    # 1. Diarize
    print(f"Diarizing {audio_path}...", file=sys.stderr)
    diarization = diarization_pipeline(audio_path)

    # 2. Transcribe full audio
    print(f"Transcribing {audio_path}...", file=sys.stderr)
    asr_result = asr_model.transcribe(
        audio_path, language="zh", fp16=torch.cuda.is_available()
    )

    # 3. Align segments: assign each ASR segment to nearest diarized speaker window
    segments = []
    for segment in asr_result["segments"]:
        start, end, text = segment["start"], segment["end"], segment["text"].strip()
        if not text:
            continue

        # Find dominant speaker in [start, end]
        best_speaker = None
        max_overlap = 0.0
        for turn, _, speaker in diarization.itertracks(yield_label=True):
            overlap = max(0, min(end, turn.end) - max(start, turn.start))
            if overlap > max_overlap:
                max_overlap = overlap
                best_speaker = speaker

        if best_speaker and max_overlap >= min_speaker_duration:
            segments.append(
                {
                    "speaker": best_speaker,
                    "text": text,
                    "start": start,
                    "end": end,
                }
            )

    return segments


def emit_jsonl(
    segments: List[Dict[str, Any]],
    session_id: str,
    timestamp_base: Optional[datetime] = None,
):
    """Emit one JSONL line per segment, then tool_use."""
    if timestamp_base is None:
        timestamp_base = datetime.now(timezone.utc)

    # Emit user messages: one per segment
    for i, seg in enumerate(segments):
        text_line = f"Speaker {seg['speaker']}: {seg['text']}"
        entry = {
            "type": "user",
            "message": {"content": [{"type": "text", "text": text_line}]},
            "timestamp": (
                timestamp_base.replace(microsecond=0) + (i * 0.1)
            ).isoformat(),  # stagger slightly for ordering
            "sessionId": session_id,
        }
        print(json.dumps(entry, ensure_ascii=False))

    # Emit assistant tool_use for summarization
    if segments:
        tool_input = {
            "segments": segments,
            "context": "Voice note from get笔记录音卡",
        }
        tool_entry = {
            "type": "assistant",
            "message": {
                "content": [
                    {
                        "type": "tool_use",
                        "id": f"t{int(time.time())}",
                        "name": "SummarizeConversation",
                        "input": tool_input,
                    }
                ]
            },
            "timestamp": timestamp_base.isoformat(),
            "sessionId": session_id,
        }
        print(json.dumps(tool_entry, ensure_ascii=False))


def main():
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument(
        "audio_path", type=str, help="Path to audio file (.m4a, .mp3, .wav)"
    )
    args = parser.parse_args()

    try:
        asr_model, diarization_pipeline = load_models()

        # Generate sessionId from file mtime
        p = Path(args.audio_path)
        mtime = datetime.fromtimestamp(p.stat().st_mtime, tz=timezone.utc)
        session_id = f"audio_{mtime.strftime('%Y%m%d_%H%M%S')}_{p.stem[:8]}"

        segments = transcribe_and_diarize(
            args.audio_path, asr_model, diarization_pipeline
        )
        emit_jsonl(segments, session_id, mtime)

    except Exception as e:
        print(f"ERROR: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()
