#!/usr/bin/env python3
"""Xiaomi MiMo ASR wrapper for Hermes command STT provider.

Usage (called by Hermes via placeholders):
  xiaomi-asr.py <input_audio> <output_txt> [model] [language]

Xiaomi ASR uses chat/completions API with input_audio format,
NOT the standard /v1/audio/transcriptions endpoint.
Only wav and mp3 are accepted — other formats are auto-converted
to 16kHz mono WAV via ffmpeg.
"""

import base64
import json
import os
import subprocess
import sys
import tempfile
import urllib.error
import urllib.request
import ssl


def main():
    if len(sys.argv) < 3:
        print(f"usage: {sys.argv[0]} <input> <output> [model] [language]", file=sys.stderr)
        sys.exit(1)

    input_path = sys.argv[1]
    output_path = sys.argv[2]
    model = sys.argv[3] if len(sys.argv) > 3 else "mimo-v2.5-asr"
    language = sys.argv[4] if len(sys.argv) > 4 else "zh"

    api_key = os.environ.get("XIAOMI_API_KEY")
    if not api_key:
        print("XIAOMI_API_KEY is not set", file=sys.stderr)
        sys.exit(1)

    base_url = os.environ.get("XIAOMI_BASE_URL", "https://api.xiaomimimo.com/v1")

    # Normalize Windows path: backslashes -> forward slashes
    input_path = input_path.replace("\\", "/")

    # Determine audio format from extension
    ext = os.path.splitext(input_path)[1].lower()
    if ext in (".wav",):
        fmt = "wav"
        audio_file = input_path
    elif ext in (".mp3",):
        fmt = "mp3"
        audio_file = input_path
    else:
        # Auto-convert to WAV via ffmpeg
        fmt = "wav"
        try:
            subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
        except (FileNotFoundError, subprocess.CalledProcessError):
            print(f"ffmpeg is required to convert {ext} to wav/mp3 for Xiaomi ASR",
                  file=sys.stderr)
            sys.exit(1)

        fd, audio_file = tempfile.mkstemp(suffix=".wav", prefix="hermes-asr-")
        os.close(fd)
        try:
            subprocess.run(
                ["ffmpeg", "-y", "-i", input_path, "-ar", "16000", "-ac", "1",
                 "-f", "wav", audio_file],
                capture_output=True, check=True
            )
        except subprocess.CalledProcessError as e:
            print(f"ffmpeg conversion failed: {e.stderr.decode(errors='replace')}",
                  file=sys.stderr)
            os.unlink(audio_file)
            sys.exit(1)

    try:
        # Read and encode audio
        with open(audio_file, "rb") as f:
            audio_b64 = base64.b64encode(f.read()).decode()

        # Build request
        payload = json.dumps({
            "model": model,
            "messages": [{
                "role": "user",
                "content": [
                    {"type": "input_audio", "input_audio": {"data": audio_b64, "format": fmt}},
                ]
            }]
        }).encode()

        req = urllib.request.Request(
            f"{base_url}/chat/completions",
            data=payload,
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
        )
        ctx = ssl.create_default_context()

        resp = urllib.request.urlopen(req, context=ctx, timeout=300)
        response_data = json.loads(resp.read().decode())

        # Extract transcript
        choices = response_data.get("choices", [])
        if not choices:
            error_info = response_data.get("error", {})
            print(error_info.get("message", "No transcription result"), file=sys.stderr)
            sys.exit(1)

        msg = choices[0].get("message", {})
        transcript = msg.get("content", "").strip()

        if not transcript:
            print("Xiaomi ASR returned empty transcript", file=sys.stderr)
            sys.exit(1)

        # Write output
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(transcript)
            f.write("\n")

        print(transcript)

    except urllib.error.HTTPError as e:
        body = e.read().decode(errors="replace")
        print(f"HTTP {e.code}: {body}", file=sys.stderr)
        sys.exit(1)
    except urllib.error.URLError as e:
        print(f"URL error: {e.reason}", file=sys.stderr)
        sys.exit(1)
    finally:
        # Clean up temp file if we created one
        if 'audio_file' in locals() and audio_file != input_path and os.path.exists(audio_file):
            os.unlink(audio_file)


if __name__ == "__main__":
    main()