"""Language Translation example.

Language Translation (LT) runs on the remote-inference pipeline and is selected
by the presence of ``lt_config`` (NOT by a model key), so AudioAttributes.
model_name is left empty. Translated audio comes back from process_frame; the
transcripts arrive on the lt_config callback.

Environment variables
---------------------
SANAS_API_KEY           (required) licence key
SANAS_STORAGE_DIR       SDK storage directory (data + logs) (default: ./storage)
SANAS_LANG_IN           source language code               (default: en-US)
SANAS_LANG_OUT          target language code               (default: es-ES)
SANAS_CONVERSATION_ID   conversation id linking sessions (default: "")
SANAS_LT_DRAIN_SECONDS  silence drain window in seconds (default: 5.0) — must
                    exceed the LT round-trip latency (~3-4 s) to capture the tail

The input and output WAV are command-line arguments (--input / --output).

Usage::

    python language_translation_example.py --input test_input.wav --output translated.wav
"""

import argparse
import os
import sys

_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _HERE)

import sanas  # noqa: E402
from helpers import PipelineWaiter, feed_and_drain, load_samples  # noqa: E402
from wav_utils import save_wav  # noqa: E402


def print_transcript(tf):
    """Pretty-print a TextFrame (translation transcript)."""
    label = (
        "Translation"
        if tf.type_ == sanas.TranscriptType.TRANSLATION
        else "Transcription"
    )
    td = tf.transcript_data_
    lang = f" lang={td.language_code_}" if td.language_code_ else ""
    print(f"\n[{label}] utterance {td.utterance_idx}{lang}")
    for kind, segs in (("complete", td.complete), ("partial", td.partial)):
        for s in segs:
            redacted = f' (redacted: "{s.redacted_text}")' if s.redacted_text else ""
            print(
                f'  [{kind}] "{s.text}"{redacted} '
                f"[{s.start}-{s.end} p={s.probability:.2f}]"
            )


def _run():
    parser = argparse.ArgumentParser(
        description="Language Translation over the Sanas SDK remote pipeline."
    )
    parser.add_argument(
        "--input",
        default="test_input.wav",
        help="input WAV path (default: %(default)s)",
    )
    parser.add_argument(
        "--output",
        default="translated.wav",
        help="output WAV path (default: %(default)s)",
    )
    args = parser.parse_args()

    api_key = os.environ.get("SANAS_API_KEY")
    if not api_key:
        sys.exit("SANAS_API_KEY is not set.")

    storage_dir = os.environ.get("SANAS_STORAGE_DIR", "./storage")
    lang_in = os.environ.get("SANAS_LANG_IN", "en-US")
    lang_out = os.environ.get("SANAS_LANG_OUT", "es-ES")
    conversation_id = os.environ.get("SANAS_CONVERSATION_ID", "conversation-id-1")

    sdk = sanas.create_sdk(sanas.InitParams(storage_dir=storage_dir))
    res = sdk.activate_api_key(api_key)
    if not res.success:
        sys.exit(f"Activation failed: {res.message}")
    print(f"[activation] OK (SDK {sdk.version})")

    samples, sample_rate, channels = load_samples(args.input)
    print(
        f"[input] {sample_rate} Hz, {channels} ch, "
        f"{len(samples) // channels} frames; translating "
        f"{lang_in} -> {lang_out}"
    )

    waiter = PipelineWaiter(on_state=lambda s: print(f"[pipeline] {s}"))
    attrs = sanas.ProcessorAttributes(
        audio_attributes=sanas.AudioAttributes(
            sampling_rate=sample_rate,
            channels=channels,
            model_name="",  # LT is selected by lt_config, not a model key
            audio_pipeline_state_notify=waiter.callback,
        ),
        lt_config=sanas.LanguageTranslationConfig(
            language_in=lang_in,
            language_out=lang_out,
            conversation_id=conversation_id,
            callback=print_transcript,
        ),
    )

    # LT runs remotely with a large (~3-4 s) latency, so after the real audio we
    # keep pushing silence for a fixed window to pull back the translated tail.
    drain_seconds = float(os.environ.get("SANAS_LT_DRAIN_SECONDS", "5.0"))

    with sdk.create_audio_processor(attrs) as proc:
        print("[pipeline] waiting for RUNNING ...")
        waiter.wait_until_running(timeout=30.0)
        result = feed_and_drain(
            proc, samples, sample_rate, channels, drain_seconds=drain_seconds
        )

    combined = result["output"]
    if combined:
        save_wav(args.output, combined, sample_rate, channels)
        print(f"\n[output] {len(combined) // 4} samples -> {args.output}")
    else:
        print("\n[output] no translated audio received.")


if __name__ == "__main__":
    _run()
