#!/usr/bin/env python3
"""
G3: Chirp Transcribe — Gọi Google Cloud Speech-to-Text V2 API
- Model: chirp (chuyên đa ngôn ngữ)
- Language: my-MM
- Features: automatic punctuation + word confidence
- Speech Adaptation với bộ từ khóa Vinaya

Requirements:
    pip install google-cloud-speech

Environment:
    export GOOGLE_APPLICATION_CREDENTIALS="/path/to/service-account.json"

Usage:
    python3 chirp_transcribe.py chunks/ -o transcripts/
    python3 chirp_transcribe.py chunks/012_chunk_000.wav -o transcripts/
"""

import argparse
import json
import os
import sys
import time
from pathlib import Path

try:
    from google.cloud.speech_v2 import SpeechClient
    from google.cloud.speech_v2.types import cloud_speech
    from google.api_core import exceptions as google_exceptions
except ImportError:
    print("❌ google-cloud-speech not installed.")
    print("   Run: pip install google-cloud-speech")
    sys.exit(1)


# ─── Vinaya Glossary for Speech Adaptation ───
# These phrases help Chirp correctly recognize key terms
VINAYA_PHRASES = [
    # General terms
    "ဦးဇင်း",           # Venerable sir (vocative)
    "တရားဟော",         # Preach Dhamma
    "တရားနာ",          # Listen to Dhamma
    "တရားတော်",        # The Dhamma (honorific)
    "စာမျက်နှာ",        # Page number
    "နာယူကောင်း",      # Suitable for listening

    # Vinaya offenses
    "အာပတ်သင့်",       # Commit an offense
    "ပါရာဇိက",         # Pārājika
    "သံဃာဒိသေသ်",      # Saṅghādisesa
    "ပါစိတ်",           # Pācittiya
    "ဒုက္ကဋ်",          # Dukkaṭa
    "ထုလ္လစ္စည်း",      # Thullaccaya

    # Sangha
    "သံဃာ",            # Sangha
    "ရဟန်း",            # Monk
    "ဝိနည်း",           # Vinaya

    # Common in Vinaya commentary
    "ဒေသနာ",            # Discourse/teaching
    "သိက္ခာပုဒ်",       # Training rule
    "မဟောကောင်း",       # Not allowed to preach
    "မနာကောင်း",        # Not allowed to listen
    "အင်္ဂါ",            # Factor/limb
    "အနာပတ္တိ",        # No offense

    # Weapons/objects from Vinaya Sekhiya rules
    "ထီး",              # Umbrella
    "လက်နက်",           # Weapon
    "ဒုတ်",             # Stick
    "လေး",              # Bow
    "မြား",             # Arrow
    "ဓား",              # Sword
    "လှံ",              # Spear
    "သေနတ်",            # Gun
]


def create_recognition_config():
    """
    Tạo config Google Cloud STT V2 tối ưu cho Myanmar Phật giáo.
    
    Khác biệt then chốt so với default:
    1. model="chirp" — model chuyên đa ngôn ngữ, không phải model mặc định
    2. enable_automatic_punctuation=True — tự động ngắt câu tiếng Myanmar
    3. enable_word_confidence=True — để debug từ nào bị nhận diện kém
    4. SpeechAdaptation — nạp từ khóa Vinaya để tăng accuracy
    """
    inline_phrase_set = cloud_speech.PhraseSet(
        phrases=[
            cloud_speech.PhraseSet.Phrase(value=phrase, boost=20.0)
            for phrase in VINAYA_PHRASES
        ]
    )
    adaptation = cloud_speech.SpeechAdaptation(
        phrase_sets=[
            cloud_speech.SpeechAdaptation.AdaptationPhraseSet(
                inline_phrase_set=inline_phrase_set
            )
        ]
    )

    config = cloud_speech.RecognitionConfig(
        language_codes=["my-MM"],
        model="chirp",  # BẮT BUỘC: không dùng model mặc định
        features=cloud_speech.RecognitionFeatures(
            enable_automatic_punctuation=True,
            enable_word_confidence=True,
        ),
        adaptation=adaptation,
    )
    return config


def transcribe_file(
    client: SpeechClient,
    project_id: str,
    audio_path: str,
    output_json_path: str
) -> bool:
    """
    Transcribe a single audio file using Chirp V2.
    
    Returns:
        True if successful, False otherwise
    """
    print(f"   🎤 Transcribing: {os.path.basename(audio_path)}")

    with open(audio_path, "rb") as f:
        audio_content = f.read()

    config = create_recognition_config()

    request = cloud_speech.RecognizeRequest(
        config=config,
        content=audio_content,
        recognizer=f"projects/{project_id}/locations/asia-southeast1/recognizers/_",
    )

    try:
        response = client.recognize(request=request)
    except google_exceptions.GoogleAPICallError as e:
        print(f"   ❌ API Error: {e}")
        return False

    # Extract results
    results = []
    for result in response.results:
        if result.alternatives:
            alt = result.alternatives[0]
            entry = {
                "transcript": alt.transcript,
                "confidence": alt.confidence,
                "words": []
            }
            for word in alt.words:
                entry["words"].append({
                    "word": word.word,
                    "start": f"{word.start_offset.total_seconds():.3f}"
                        if word.start_offset else "0.000",
                    "end": f"{word.end_offset.total_seconds():.3f}"
                        if word.end_offset else "0.000",
                    "confidence": word.confidence,
                })
            results.append(entry)

    # Save raw JSON
    output = {
        "audio_file": audio_path,
        "model": "chirp",
        "language": "my-MM",
        "results": results
    }

    with open(output_json_path, "w", encoding="utf-8") as f:
        json.dump(output, f, ensure_ascii=False, indent=2)

    total_chars = sum(len(r["transcript"]) for r in results)
    print(f"   ✅ {len(results)} segment(s), {total_chars} chars → {output_json_path}")
    return True


def transcribe_directory(
    client: SpeechClient,
    project_id: str,
    input_dir: str,
    output_dir: str,
    retry_failed: bool = True,
    max_retries: int = 3
):
    """
    Transcribe all .wav files in a directory.
    Supports retry for failed chunks.
    """
    wav_files = sorted(Path(input_dir).glob("*.wav"))
    if not wav_files:
        print(f"❌ No .wav files found in {input_dir}")
        return

    os.makedirs(output_dir, exist_ok=True)
    print(f"\n🎯 Transcribing {len(wav_files)} file(s) from {input_dir}/")
    print(f"   Model: chirp | Language: my-MM | Phrases: {len(VINAYA_PHRASES)}")
    print()

    failed = []
    success = 0

    for wav_path in wav_files:
        json_name = wav_path.stem + "_raw.json"
        json_path = os.path.join(output_dir, json_name)

        # Skip if already processed
        if os.path.exists(json_path):
            print(f"   ⏭️  Skipping {wav_path.name} (already processed)")
            success += 1
            continue

        ok = False
        for attempt in range(max_retries):
            if attempt > 0:
                wait = 2 ** attempt
                print(f"   🔄 Retry {attempt}/{max_retries} after {wait}s...")
                time.sleep(wait)

            ok = transcribe_file(client, project_id, str(wav_path), json_path)
            if ok:
                success += 1
                break

        if not ok:
            failed.append(wav_path.name)

    print(f"\n{'='*50}")
    print(f"📊 Summary: {success}/{len(wav_files)} succeeded")
    if failed:
        print(f"❌ Failed: {', '.join(failed)}")
    return failed


def main():
    parser = argparse.ArgumentParser(
        description="G3: Chirp Transcribe — Google Cloud STT V2 for Myanmar"
    )
    parser.add_argument("input", help="Input .wav file or directory of .wav files")
    parser.add_argument("-o", "--output-dir", default="transcripts/",
                        help="Output directory for raw JSON transcripts")
    parser.add_argument("-p", "--project-id", default=None,
                        help="Google Cloud project ID (default: from env)")
    parser.add_argument("--max-retries", type=int, default=3,
                        help="Max API retries per chunk (default: 3)")
    parser.add_argument("--no-retry", action="store_true",
                        help="Skip retry for failed chunks")
    args = parser.parse_args()

    # Get project ID from env or args
    project_id = args.project_id or os.environ.get("GOOGLE_CLOUD_PROJECT")
    if not project_id:
        # Try to read from credentials file
        creds_path = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", "")
        if creds_path and os.path.exists(creds_path):
            with open(creds_path) as f:
                creds = json.load(f)
                project_id = creds.get("project_id")

    if not project_id:
        print("❌ Project ID not found. Set GOOGLE_CLOUD_PROJECT or "
              "GOOGLE_APPLICATION_CREDENTIALS")
        sys.exit(1)

    print(f"🔑 Project: {project_id} | Location: asia-southeast1")

    client = SpeechClient()

    if os.path.isdir(args.input):
        transcribe_directory(
            client, project_id, args.input, args.output_dir,
            retry_failed=not args.no_retry,
            max_retries=args.max_retries
        )
    else:
        json_name = Path(args.input).stem + "_raw.json"
        json_path = os.path.join(args.output_dir, json_name)
        os.makedirs(args.output_dir, exist_ok=True)
        transcribe_file(client, project_id, args.input, json_path)


if __name__ == "__main__":
    main()