#!/usr/bin/env python3
"""
Smart STT Pipeline v3
T-one (fast) + GigaAM (accurate) + Qwen3.5-4B (punctuation)
Auto-selects best STT based on confidence/silence detection.

Configuration: all sensitive data loaded from .env file.
"""

import os
import tempfile
import time
import subprocess
import re

from dotenv import load_dotenv
from fastapi import FastAPI, File, UploadFile, HTTPException
from pydantic import BaseModel
import requests

# Load .env (skip if not found)
load_dotenv()

# === Configuration from .env ===
TONE_URL = os.getenv("TONE_URL", "http://localhost:8003/transcribe")
GIGAAM_URL = os.getenv("GIGAAM_URL", "http://localhost:8001/transcribe")
QWEN_URL = os.getenv("QWEN_URL", "http://178.17.143.46:8080/v1/chat/completions")
QWEN_URL_FALLBACK = os.getenv("QWEN_URL_FALLBACK", "http://192.168.10.101:8888/v1/chat/completions")
QWEN_API_KEY = os.getenv("QWEN_API_KEY", "")
QWEN_MODEL = os.getenv("QWEN_MODEL", "Qwen3.5-4B-Q5_K_M.gguf")
QWEN_MODEL_FALLBACK = os.getenv("QWEN_MODEL_FALLBACK", "Qwen3.5-4B-MTP-GGUF")
MIN_TEXT_LENGTH = int(os.getenv("MIN_TEXT_LENGTH", "5"))
SHORT_AUDIO_THRESHOLD = int(os.getenv("SHORT_AUDIO_THRESHOLD", "10"))

app = FastAPI(title="Smart STT Pipeline v3", version="3.0.0")


class STTResponse(BaseModel):
    original: str
    punctuated: str
    duration_sec: float
    stt_time_sec: float
    punct_time_sec: float
    server_used: str
    stt_engine: str  # "t-one" or "gigaam"


def convert_to_wav(audio_bytes: bytes) -> str:
    """Convert any audio to 16kHz mono WAV"""
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
        tmp_path = tmp.name
    cmd = [
        "ffmpeg", "-i", "pipe:0",
        "-ar", "16000", "-ac", "1", "-c:a", "pcm_s16le",
        "-y", tmp_path
    ]
    subprocess.run(cmd, input=audio_bytes, capture_output=True)
    return tmp_path


def transcribe_tone(wav_path: str) -> dict:
    """Transcribe using T-one, returns dict with text and validity"""
    with open(wav_path, "rb") as f:
        files = {"file": ("audio.wav", f, "audio/wav")}
        resp = requests.post(TONE_URL, files=files, timeout=30)
    data = resp.json()
    if "detail" in data:
        return {"text": "", "valid": False, "error": data["detail"]}

    text = data.get("text", "[]")
    phrases = re.findall(r"TextPhrase\(text='([^']*)'", text)
    result_text = " ".join(phrases)

    # Check if result is meaningful
    valid = len(result_text.strip()) >= MIN_TEXT_LENGTH

    return {
        "text": result_text,
        "valid": valid,
        "stt_time": data.get("stt_time_sec", 0)
    }


def transcribe_gigaam(ogg_path: str) -> dict:
    """Transcribe using GigaAM, returns dict with full response"""
    with open(ogg_path, "rb") as f:
        files = {"file": ("audio.ogg", f, "audio/ogg")}
        resp = requests.post(GIGAAM_URL, files=files, timeout=60)
    return resp.json()


def add_punctuation(text: str, server_url: str, api_key: str, model: str) -> str:
    """Add punctuation using Qwen3.5-4B"""
    payload = {
        "model": model,
        "messages": [
            {
                "role": "system",
                "content": "You are a punctuation assistant. Add punctuation to Russian text. Only output the punctuated text, no explanations."
            },
            {"role": "user", "content": text}
        ],
        "max_tokens": 500,
        "temperature": 0.1
    }
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    resp = requests.post(server_url, json=payload, headers=headers, timeout=60)
    resp.raise_for_status()
    return resp.json()["choices"][0]["message"]["content"]


@app.post("/transcribe")
async def transcribe(file: UploadFile = File(...)):
    """Smart transcription: T-one first, fallback to GigaAM if needed"""
    audio_bytes = await file.read()

    # Get audio duration from file
    ogg_path = tempfile.mktemp(suffix=".ogg")
    with open(ogg_path, "wb") as f:
        f.write(audio_bytes)

    cmd = ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", ogg_path]
    proc = subprocess.run(cmd, capture_output=True, text=True)
    duration = float(proc.stdout.strip())

    # Convert to WAV for T-one
    wav_path = convert_to_wav(audio_bytes)

    try:
        # Step 1: Choose STT engine based on duration
        if duration < SHORT_AUDIO_THRESHOLD:
            # Short audio: use GigaAM (more accurate)
            giga_result = transcribe_gigaam(ogg_path)

            text = giga_result.get("original", "")
            stt_engine = "gigaam"
            stt_time = giga_result.get("stt_time_sec", 0)
        else:
            # Long audio: try T-one first (fast)
            stt_start = time.time()
            tone_result = transcribe_tone(wav_path)
            stt_time = time.time() - stt_start

            if tone_result["valid"]:
                text = tone_result["text"]
                stt_engine = "t-one"
            else:
                # T-one failed, fallback to GigaAM
                giga_result = transcribe_gigaam(ogg_path)
                text = giga_result.get("original", "")
                stt_engine = "gigaam"
                stt_time = giga_result.get("stt_time_sec", 0)

        # Step 2: Punctuation
        punct_start = time.time()
        server_used = "178.17.143.46:8080"
        try:
            punctuated = add_punctuation(text, QWEN_URL, QWEN_API_KEY, QWEN_MODEL)
        except Exception:
            server_used = "192.168.10.101:8888"
            punctuated = add_punctuation(text, QWEN_URL_FALLBACK, QWEN_API_KEY, QWEN_MODEL_FALLBACK)
        punct_time = time.time() - punct_start

        return {
            "original": text,
            "punctuated": punctuated,
            "duration_sec": round(duration, 1),
            "stt_time_sec": round(stt_time, 2),
            "punct_time_sec": round(punct_time, 2),
            "server_used": server_used,
            "stt_engine": stt_engine
        }
    finally:
        os.unlink(wav_path)


@app.get("/health")
async def health():
    return {"status": "ok", "service": "stt-pipeline-v3"}