#!/usr/bin/env python3
"""Translate a Persian transcript JSON to English via Claude, batched, Sufi-aware.

Reads <input>.fa.json (produced by transcribe.py).
Writes <input>.en.json with the same segment timestamps but English text.

Batches segments to keep token cost down + give Claude paragraph context (better
translation than line-at-a-time). Sufi-aware prompt preserves poetic register +
technical religious vocabulary.

Usage:
    python translate.py audio/<file>.mp3                      # uses <file>.mp3.fa.json
    python translate.py audio/<file>.mp3 --batch 20           # segments per Claude call
    python translate.py audio/<file>.mp3 --force              # re-translate
"""
import argparse
import json
import os
import re
import sys
from pathlib import Path

import anthropic

CLAUDE_MODEL = "claude-sonnet-4-6"

SYSTEM_PROMPT = """You are a translator working on Persian (Farsi) Sufi lecture transcripts \
for Cameron, a Persian-American student of the Maktab Tariqat Oveyssi tradition. \
He's an intermediate Persian speaker; the translations are for personal study + \
listening alongside the original audio.

Translate each numbered Persian segment to clear, natural English. CRITICAL constraints:

1. ONE-TO-ONE mapping. Each numbered Persian segment gets exactly one numbered English \
translation. Do not merge segments, do not split, do not skip. This is for subtitle alignment.

2. Preserve register. Sufi material is poetic + reverent. Don't flatten into corporate \
English. If the Persian uses *delam* (heart), keep "heart," not "feeling." If it uses \
*'eshq* (divine love), translate as "love" with capital L *only* if context implies the \
divine; otherwise lowercase. Keep technical Sufi terms in transliteration when no clean \
English equivalent exists (e.g., *dhikr*, *fana*, *Pir*, *tariqat*).

3. Don't editorialize or add interpretation. If a sentence is fragmented in Persian (often \
the case with extemporaneous lectures), the English can be fragmented too. Don't smooth.

4. If the Persian is unclear (Whisper error, garbled), output the best-guess English in \
square brackets, e.g., "[unclear: possibly about prayer]".

5. Keep it concise. English subtitle lines should be roughly the same reading time as the \
Persian segment's playback time.

Output format — EXACTLY:
1. <english translation>
2. <english translation>
3. <english translation>
...

No preamble, no commentary, no explanation. Just the numbered translations."""


def parse_batch_output(text: str, expected_count: int) -> list[str]:
    """Parse Claude's '1. ...\n2. ...' output back into a list."""
    out = [""] * expected_count
    for line in text.splitlines():
        line = line.strip()
        m = re.match(r"^(\d+)[.\):]\s*(.*)", line)
        if not m:
            continue
        idx = int(m.group(1)) - 1
        if 0 <= idx < expected_count:
            out[idx] = m.group(2).strip()
    return out


def translate_batch(client: anthropic.Anthropic, segments_subset: list[dict]) -> list[str]:
    numbered = "\n".join(f"{i+1}. {s['text']}" for i, s in enumerate(segments_subset))
    resp = client.messages.create(
        model=CLAUDE_MODEL,
        max_tokens=4096,
        system=SYSTEM_PROMPT,
        messages=[{"role": "user", "content": f"Translate these {len(segments_subset)} Persian segments:\n\n{numbered}"}],
    )
    return parse_batch_output(resp.content[0].text, len(segments_subset))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("audio_file", type=Path)
    ap.add_argument("--batch", type=int, default=20)
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()

    fa_path = args.audio_file.with_suffix(args.audio_file.suffix + ".fa.json")
    en_path = args.audio_file.with_suffix(args.audio_file.suffix + ".en.json")
    if not fa_path.exists():
        sys.exit(f"Persian transcript not found: {fa_path} — run transcribe.py first")
    if en_path.exists() and not args.force:
        print(f"already translated → {en_path} (use --force to redo)")
        return

    data = json.loads(fa_path.read_text())
    segments = data["segments"]
    print(f"translating {len(segments)} segments in batches of {args.batch}…")

    client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    translations = []
    for i in range(0, len(segments), args.batch):
        batch = segments[i:i + args.batch]
        print(f"  batch {i//args.batch + 1} ({i+1}-{i+len(batch)})...", flush=True)
        ens = translate_batch(client, batch)
        translations.extend(ens)

    if len(translations) != len(segments):
        sys.exit(f"length mismatch: {len(translations)} vs {len(segments)}")

    out_segments = [
        {**seg, "english": en}
        for seg, en in zip(segments, translations)
    ]
    en_path.write_text(json.dumps({
        "source": data["source"],
        "model": CLAUDE_MODEL,
        "segments": out_segments,
    }, ensure_ascii=False, indent=2))
    en_path.chmod(0o600)
    print(f"\n✓ {len(out_segments)} translations → {en_path}")


if __name__ == "__main__":
    main()