#!/usr/bin/env python3
"""
G4: Gemini Post-Processing — Hiệu đính ngữ cảnh transcript Myanmar
- Model: Gemini 3 Flash Preview qua OpenRouter
- Prompt chuyên gia Phật giáo Myanmar
- Khôi phục trợ từ, sửa đồng âm, phân đoạn mạch lạc

Requirements:
    pip install openai

Environment:
    export OPENROUTER_API_KEY="your-openrouter-key"

Usage:
    python3 gemini_postprocess.py transcripts/ -o transcripts/
    python3 gemini_postprocess.py transcripts/012_chunk_000_raw.json -o transcripts/
"""

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

try:
    from openai import OpenAI
except ImportError:
    print("❌ openai SDK not installed.")
    print("   Run: pip install openai")
    sys.exit(1)


# ─── System Prompt cho Gemini (qua OpenRouter) ───
# Đây là yếu tố quyết định chất lượng đầu ra.
# Thiết kế để Gemini hoạt động như chuyên gia hiệu đính Phật giáo Myanmar.

SYSTEM_PROMPT = """Bạn là một chuyên gia hiệu đính tài liệu Phật giáo Myanmar bậc cao. 
Bạn sẽ nhận được một bản transcript thô từ hệ thống nhận diện giọng nói tự động (STT) 
từ âm thanh hội trường — nơi có tiếng vang, giọng tụng đều đều, và nhiều trợ từ cổ.

**Nhiệm vụ của bạn:**

1. **Đọc toàn bộ văn bản thô** để nắm bắt chủ đề chính. 
   Ví dụ: Đang giảng về Luật tạng (Vinaya), trang bao nhiêu, các quy định về việc 
   không thuyết pháp cho người cầm che ô/cầm vũ khí, lỗi phạm giới.

2. **Điền lại các trợ từ bị mất** do STT bỏ qua — các từ không mang nghĩa từ vựng 
   nhưng quan trọng cho ngữ pháp Myanmar chuẩn:
   - ၏ (kết thúc câu trang trọng)
   - သာလျှင် (chỉ khi)
   - ပေါ့နော်, နော်, ပေါ့ (trợ từ khẳng định cuối câu)
   - လျှင် (nếu)
   - ပြီးတော့ (sau đó)
   - ကို (tân ngữ)
   - မှ (mới)

3. **Sửa các lỗi nhận diện từ đồng âm** do STT sai:
   - Ví dụ: Máy trích xuất 'ထိတောင်' → bạn phải sửa thành 'ထီးတောင်ဝှေး' 
     nếu ngữ cảnh đang nói về cây gậy chống
   - 'ဒုတ်' vs 'ထုတ်' → dựa vào việc vị sư đang nói về gậy hay hành động rút ra
   - 'ပေါင်' vs 'ပောင်' → dựa vào ngữ cảnh
   - 'လက်အုပ်' vs 'လုပ်ချင်' → lễ phép chắp tay vs muốn làm gì

4. **Tuyệt đối không tự ý thêm thắt** nội dung mới không có trong mạch ý chính 
   của đoạn văn thô. Chỉ sửa lỗi, không bịa ra nội dung.

5. **Trả về kết quả** là văn bản tiếng Myanmar sạch, phân đoạn mạch lạc 
   theo ý giảng, không thêm bất kỳ chú thích hay meta-data nào.

**Định dạng output:**
Chỉ trả về văn bản tiếng Myanmar đã hiệu đính, 
phân đoạn bằng dòng trống giữa các ý chính.
KHÔNG thêm tiêu đề, giải thích, hay bất kỳ nội dung nào không phải transcript."""


# OpenRouter API endpoint
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"
DEFAULT_MODEL = "google/gemini-3-flash-preview"


def extract_raw_text(json_path: str) -> str:
    """
    Extract concatenated transcript text from Chirp raw JSON.
    Handles multiple Chirp JSON structures:
    - Structure A: {"results": [{"transcript": "..."}]}  (long-form V2)
    - Structure B: {"segments": [{"transcript": "..."}]}  (batch V2)
    - Structure C: {"transcript": "..."}  (short-form)
    """
    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    texts = []

    # Structure A: results array
    for result in data.get("results", []):
        text = result.get("transcript", "").strip()
        if text:
            texts.append(text)

    # Structure B: segments array
    for seg in data.get("segments", []):
        text = seg.get("transcript", "").strip()
        if text:
            texts.append(text)

    # Structure C: direct transcript field
    direct = data.get("transcript", "").strip()
    if direct:
        texts.append(direct)

    return "\n".join(texts)


def postprocess_with_gemini(
    raw_text: str,
    api_key: str,
    chunk_name: str,
    model: str = DEFAULT_MODEL
) -> str:
    """
    Send raw transcript to Gemini via OpenRouter for contextual post-processing.
    
    Args:
        raw_text: Raw transcript from Chirp
        api_key: OpenRouter API key
        chunk_name: Name of chunk for logging
        model: Model ID (OpenRouter format, e.g. google/gemini-3-flash-preview)
    
    Returns:
        Cleaned Myanmar text
    """
    print(f"   🤖 Post-processing with {model}...")

    client = OpenAI(
        api_key=api_key,
        base_url=OPENROUTER_BASE_URL,
    )

    user_prompt = f"""Đây là transcript thô từ bài giảng Phật giáo tiếng Myanmar (đoạn: {chunk_name}).

Hãy hiệu đính theo đúng hướng dẫn hệ thống:

--- BEGIN RAW TRANSCRIPT ---
{raw_text}
--- END RAW TRANSCRIPT ---

Trả về CHỈ văn bản tiếng Myanmar đã hiệu đính, không thêm gì khác."""

    try:
        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            temperature=0.2,       # Low temp → giữ nguyên ý, ít hallucinate
            max_tokens=8192,
            extra_headers={
                "HTTP-Referer": "https://github.com/openclaw/pipeline",
                "X-Title": "Myanmar Buddhist Transcript Post-Processor",
            },
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        print(f"   ❌ OpenRouter API error: {e}")
        return ""


def process_file(
    api_key: str,
    json_path: str,
    output_txt_path: str,
    model: str = DEFAULT_MODEL,
    max_retries: int = 3
) -> bool:
    """
    Process a single raw JSON → clean TXT via Gemini.
    
    Returns:
        True if successful
    """
    chunk_name = Path(json_path).stem.replace("_raw", "")
    print(f"   📄 Processing: {os.path.basename(json_path)}")

    raw_text = extract_raw_text(json_path)
    if not raw_text.strip():
        print(f"   ⚠️  Empty transcript, skipping")
        return False

    print(f"   📊 Raw: {len(raw_text)} chars")
    # Quick sanity: if too short, likely failed transcription
    if len(raw_text) < 10:
        print(f"   ⚠️  Transcript too short ({len(raw_text)} chars), skipping")
        return 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)

        clean_text = postprocess_with_gemini(raw_text, api_key, chunk_name, model)

        if clean_text and len(clean_text) > 10:
            with open(output_txt_path, "w", encoding="utf-8") as f:
                f.write(clean_text)
            print(f"   ✅ Clean: {len(clean_text)} chars → {output_txt_path}")
            return True

    print(f"   ❌ Failed after {max_retries} retries")
    return False


def process_directory(
    api_key: str,
    input_dir: str,
    output_dir: str,
    model: str = DEFAULT_MODEL,
    max_retries: int = 3
):
    """
    Process all *_raw.json files in a directory.
    """
    json_files = sorted(Path(input_dir).glob("*_raw.json"))
    if not json_files:
        print(f"❌ No *_raw.json files found in {input_dir}")
        return

    os.makedirs(output_dir, exist_ok=True)
    print(f"\n🧠 Gemini Post-Processing {len(json_files)} file(s)")
    print(f"   Model: {model} | Temperature: 0.2 | Provider: OpenRouter")
    print()

    success = 0
    failed = []

    for json_path in json_files:
        txt_name = json_path.stem.replace("_raw", "_clean") + ".txt"
        txt_path = os.path.join(output_dir, txt_name)

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

        ok = process_file(api_key, str(json_path), txt_path, model, max_retries)
        if ok:
            success += 1
        else:
            failed.append(json_path.name)

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


def main():
    parser = argparse.ArgumentParser(
        description="G4: Gemini Post-Processing for Myanmar Transcripts (via OpenRouter)"
    )
    parser.add_argument("input", help="Input _raw.json file or directory")
    parser.add_argument("-o", "--output-dir", default=None,
                        help="Output directory (default: same as input)")
    parser.add_argument("-k", "--api-key", default=None,
                        help="OpenRouter API key (default: from OPENROUTER_API_KEY env)")
    parser.add_argument("-m", "--model", default=DEFAULT_MODEL,
                        help=f"Model via OpenRouter (default: {DEFAULT_MODEL})")
    parser.add_argument("--max-retries", type=int, default=3,
                        help="Max API retries per chunk (default: 3)")
    args = parser.parse_args()

    api_key = args.api_key or os.environ.get("OPENROUTER_API_KEY")
    if not api_key:
        print("❌ OpenRouter API key required. Set OPENROUTER_API_KEY or use -k")
        sys.exit(1)

    output_dir = args.output_dir or (
        args.input if os.path.isdir(args.input) else os.path.dirname(args.input)
    )

    if os.path.isdir(args.input):
        process_directory(api_key, args.input, output_dir, args.model, args.max_retries)
    else:
        json_path = args.input
        txt_name = Path(json_path).stem.replace("_raw", "_clean") + ".txt"
        txt_path = os.path.join(output_dir, txt_name)
        os.makedirs(output_dir, exist_ok=True)
        process_file(api_key, json_path, txt_path, args.model, args.max_retries)


if __name__ == "__main__":
    main()