import argparse
import json
import sys
import time
import unicodedata
from contextlib import redirect_stdout

from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess


def clean_text(value):
    if not value:
        return ""
    text = rich_transcription_postprocess(value)
    text = unicodedata.normalize("NFKC", text)
    return " ".join(text.split()).strip()


def emit(payload):
    print(json.dumps(payload, ensure_ascii=False), flush=True)


def transcribe(model, audio_path):
    started = time.perf_counter()
    # FunASR may write progress messages to stdout.  Stdout is reserved for
    # the one-line JSON protocol consumed by the desktop client.
    with redirect_stdout(sys.stderr):
        result = model.generate(
            input=audio_path,
            language="zh",
            use_itn=True,
            batch_size_s=60,
        )

    raw = ""
    if result:
        first = result[0] if isinstance(result, list) else result
        if isinstance(first, dict):
            raw = str(first.get("text", ""))
        else:
            raw = str(first)
    return clean_text(raw), round((time.perf_counter() - started) * 1000)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-dir", required=True)
    parser.add_argument("--vad-model-dir", required=True)
    parser.add_argument("--device", default="cpu")
    args = parser.parse_args()

    # The model is loaded once, then the process accepts sequential requests
    # over stdin.  This eliminates the per-10-second-clip model boot cost.
    with redirect_stdout(sys.stderr):
        model = AutoModel(
            model=args.model_dir,
            vad_model=args.vad_model_dir,
            vad_kwargs={"max_single_segment_time": 30000},
            device=args.device,
            disable_update=True,
        )
    emit({"ready": True, "engine": "sensevoice-small"})

    for line in sys.stdin:
        # Some Windows hosts prepend an UTF-8 BOM to redirected stdin.
        line = line.lstrip("\ufeff").removeprefix("ï»¿")
        if not line.strip():
            continue
        request_id = ""
        try:
            request = json.loads(line)
            request_id = str(request.get("id", ""))
            audio_path = str(request.get("audio", "")).strip()
            if not audio_path:
                raise ValueError("missing audio path")
            text, latency_ms = transcribe(model, audio_path)
            emit({
                "id": request_id,
                "ok": True,
                "text": text,
                "latency_ms": latency_ms,
                "language": "zh",
                "engine": "sensevoice-small",
            })
        except Exception as exc:
            emit({
                "id": request_id,
                "ok": False,
                "error": str(exc),
                "engine": "sensevoice-small",
            })


if __name__ == "__main__":
    sys.stdout.reconfigure(encoding="utf-8")
    sys.stderr.reconfigure(encoding="utf-8")
    main()
