#!/usr/bin/env python3
"""OpenAI-router STT adapter for the persistent, CPU-only Qwen3-ASR server."""

import json
import logging
import os
import subprocess
import tempfile
import time
import urllib.request
import uuid
from email import policy
from email.parser import BytesParser
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer


HOST = os.environ.get("QWEN_ASR_HOST", "0.0.0.0")
PORT = int(os.environ.get("QWEN_ASR_PORT", "8084"))
SERVER_URL = os.environ.get("QWEN_ASR_SERVER_URL", "http://qwen-asr:8080").rstrip("/")
LANGUAGE = os.environ.get("QWEN_ASR_LANGUAGE", "de")
MAX_BODY_BYTES = 25 * 1024 * 1024

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("qwen-asr-worker")


def clean_transcript(value: str) -> str:
    """llama.cpp may include a Qwen task marker before the spoken words."""
    if "<asr_text>" in value:
        value = value.split("<asr_text>", 1)[1]
    return value.replace("<|endoftext|>", "").strip()


def transcribe(audio: bytes, filename: str, language: str) -> dict:
    suffix = os.path.splitext(filename)[1].lower() or ".wav"
    with tempfile.TemporaryDirectory(prefix="qwen_asr_") as directory:
        source = os.path.join(directory, "input" + suffix)
        wav = os.path.join(directory, "audio.wav")
        with open(source, "wb") as handle:
            handle.write(audio)
        result = subprocess.run(
            ["ffmpeg", "-nostdin", "-hide_banner", "-loglevel", "error", "-y",
             "-i", source, "-ac", "1", "-ar", "16000", "-c:a", "pcm_s16le", wav],
            capture_output=True, text=True, timeout=30,
        )
        if result.returncode:
            raise ValueError("Audio konnte nicht gelesen werden: " + result.stderr[-300:])
        with open(wav, "rb") as handle:
            pcm = handle.read()

    boundary = "athena-qwen-asr-" + uuid.uuid4().hex
    body = b"".join([
        f"--{boundary}\r\n".encode(),
        b'Content-Disposition: form-data; name="file"; filename="audio.wav"\r\n',
        b"Content-Type: audio/wav\r\n\r\n", pcm, b"\r\n",
        f"--{boundary}\r\n".encode(),
        b'Content-Disposition: form-data; name="model"\r\n\r\n',
        b"qwen3-asr-0.6b\r\n",
        f"--{boundary}\r\n".encode(),
        b'Content-Disposition: form-data; name="language"\r\n\r\n',
        language.encode(), b"\r\n",
        f"--{boundary}--\r\n".encode(),
    ])
    request = urllib.request.Request(
        SERVER_URL + "/v1/audio/transcriptions", data=body,
        headers={"Content-Type": f"multipart/form-data; boundary={boundary}"},
        method="POST",
    )
    started = time.monotonic()
    with urllib.request.urlopen(request, timeout=60) as response:
        payload = json.loads(response.read())
    if not isinstance(payload, dict) or not isinstance(payload.get("text"), str):
        raise RuntimeError("Qwen3-ASR returned no transcription")
    text = clean_transcript(payload["text"])
    elapsed = int((time.monotonic() - started) * 1000)
    log.info("Qwen3-ASR transcribed %d characters in %d ms", len(text), elapsed)
    return {"text": text, "language": language, "duration_ms": elapsed,
            "engine": "qwen3-asr-0.6b"}


def parse_audio(body: bytes, content_type: str) -> tuple[bytes, str, str]:
    if "multipart/form-data" not in content_type.lower():
        return body, "audio.wav", LANGUAGE
    message = BytesParser(policy=policy.default).parsebytes(
        b"MIME-Version: 1.0\r\nContent-Type: " + content_type.encode() +
        b"\r\n\r\n" + body
    )
    if not message.is_multipart():
        raise ValueError("Invalid multipart upload")
    audio = b""
    filename = "audio.wav"
    language = LANGUAGE
    for part in message.iter_parts():
        name = part.get_param("name", header="content-disposition")
        if name == "file":
            audio = part.get_payload(decode=True) or b""
            filename = os.path.basename(part.get_filename() or filename)
        elif name == "language":
            language = (part.get_payload(decode=True) or b"").decode("utf-8").strip()
    return audio, filename, language if language and language != "auto" else LANGUAGE


class Handler(BaseHTTPRequestHandler):
    def send_json(self, status: int, data: dict) -> None:
        body = json.dumps(data, ensure_ascii=False).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self) -> None:
        if self.path != "/status":
            self.send_json(404, {"error": "not found"})
            return
        try:
            with urllib.request.urlopen(SERVER_URL + "/health", timeout=2) as response:
                ready = response.status == 200
        except Exception:
            ready = False
        self.send_json(200, {"ready": ready, "model": "qwen3-asr-0.6b",
                             "engine": "qwen3-asr", "language": LANGUAGE})

    def do_POST(self) -> None:
        if self.path != "/transcribe":
            self.send_json(404, {"error": "not found"})
            return
        try:
            size = int(self.headers.get("Content-Length", "0"))
            if not 0 < size <= MAX_BODY_BYTES:
                self.send_json(413, {"error": "Invalid audio size"})
                return
            audio, filename, language = parse_audio(
                self.rfile.read(size), self.headers.get("Content-Type", "")
            )
            if not audio:
                raise ValueError("Missing audio file")
            self.send_json(200, transcribe(audio, filename, language))
        except ValueError as exc:
            self.send_json(400, {"error": str(exc)})
        except Exception:
            log.exception("Transcription failed")
            self.send_json(503, {"error": "Qwen3-ASR unavailable"})


if __name__ == "__main__":
    ThreadingHTTPServer((HOST, PORT), Handler).serve_forever()