import os
import re
import time
import base64
import asyncio
import tempfile
import traceback
import subprocess
import httpx
from fastapi import FastAPI, UploadFile, File, Form, Header, HTTPException, Request, Response
from fastapi.responses import JSONResponse, FileResponse

VERSION = "1.1.0"

app = FastAPI(title="TypeWhisper-Nhi")

# Your workspace-specific Singapore compatible-mode endpoint
ALIBABA_BASE_URL = "https://ws-fnt3psiwu1morhou.ap-southeast-1.maas.aliyuncs.com/compatible-mode/v1"

# Directory to temporarily host audio files for ASR models
HOSTED_DIR = os.path.join(tempfile.gettempdir(), "qwen_hosted")
os.makedirs(HOSTED_DIR, exist_ok=True)

LOCAL_MODELS = [
    {"id": "local-zipformer-vi-68m", "object": "model", "owned_by": "local"},
    {"id": "local-zipformer-vi-68m-punc", "object": "model", "owned_by": "local"},
    {"id": "local-zipformer-vi-68m-punc-diar", "object": "model", "owned_by": "local"},
    {"id": "local-zipformer-vi-68m-punc-diar-spk2", "object": "model", "owned_by": "local"},
    {"id": "local-zipformer-vi-68m-punc-diar-spk3", "object": "model", "owned_by": "local"},
    {"id": "local-zipformer-vi-68m-punc-diar-spk4", "object": "model", "owned_by": "local"},
    {"id": "local-zipformer-vi-68m-punc-diar-spk5", "object": "model", "owned_by": "local"}
]

GEMINI_MODELS = [
    {"id": "gemini-2.5-flash", "object": "model", "owned_by": "google"},
    {"id": "gemini-2.5-pro", "object": "model", "owned_by": "google"},
    {"id": "gemini-2.0-flash", "object": "model", "owned_by": "google"},
    {"id": "gemini-1.5-flash", "object": "model", "owned_by": "google"},
    {"id": "gemini-1.5-pro", "object": "model", "owned_by": "google"}
]

def save_last_used_model(model_type: str, model_name: str):
    try:
        config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "last_used_models.json")
        data = {}
        if os.path.exists(config_path):
            with open(config_path, "r", encoding="utf-8") as f:
                data = json.load(f)
        data[model_type] = model_name
        with open(config_path, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        print(f"Saved last used {model_type} model: '{model_name}'")
    except Exception as e:
        print(f"Error saving last used model: {e}")

@app.get("/files/{filename}")
async def get_hosted_file(filename: str):
    file_path = os.path.join(HOSTED_DIR, filename)
    if os.path.exists(file_path):
        return FileResponse(file_path)
    raise HTTPException(status_code=404, detail="File not found")

def compress_audio(input_bytes: bytes, file_format: str) -> tuple[bytes, str]:
    """
    Compresses incoming audio to a low-bitrate mono MP3 (16kHz)
    using ffmpeg to ensure it fits well under Alibaba's limit.
    """
    # Create temporary files for input and output
    with tempfile.NamedTemporaryFile(suffix=f".{file_format}", delete=False) as temp_in:
        temp_in.write(input_bytes)
        temp_in_path = temp_in.name
        
    temp_out_path = tempfile.mktemp(suffix=".mp3")
    
    try:
        # Run ffmpeg to convert to 16kHz mono MP3 at 24kbps
        # This yields extremely high quality for ASR but is tiny (~180KB per minute of audio)
        ffmpeg_bin = "./ffmpeg" if os.path.exists("./ffmpeg") else "ffmpeg"
        cmd = [
            ffmpeg_bin, "-y",
            "-i", temp_in_path,
            "-ar", "16000",      # 16kHz sampling rate
            "-ac", "1",          # Mono channel
            "-b:a", "96k",       # 96kbps bitrate (preserves vocal details for speaker diarization)
            temp_out_path
        ]
        
        # Run silently
        subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
        
        # Read the compressed file
        with open(temp_out_path, "rb") as f:
            compressed_bytes = f.read()
            
        print(f"Audio compressed from {len(input_bytes) / 1024 / 1024:.2f} MB down to {len(compressed_bytes) / 1024 / 1024:.2f} MB")
        return compressed_bytes, "mp3"
        
    except Exception as e:
        print("FFmpeg compression failed, falling back to original audio:")
        traceback.print_exc()
        return input_bytes, file_format
        
    finally:
        # Clean up temporary files safely
        try:
            if os.path.exists(temp_in_path):
                os.remove(temp_in_path)
            if os.path.exists(temp_out_path):
                os.remove(temp_out_path)
        except Exception:
            pass

def convert_to_local_wav(input_bytes: bytes, file_format: str) -> str:
    """
    Converts incoming audio bytes to a standard 16kHz mono WAV file
    using ffmpeg to ensure compatibility with local Zipformer ASR.
    Returns the path to the temporary WAV file.
    """
    with tempfile.NamedTemporaryFile(suffix=f".{file_format}", delete=False) as temp_in:
        temp_in.write(input_bytes)
        temp_in_path = temp_in.name
        
    temp_out_path = tempfile.mktemp(suffix=".wav")
    
    try:
        ffmpeg_bin = "./ffmpeg" if os.path.exists("./ffmpeg") else "ffmpeg"
        cmd = [
            ffmpeg_bin, "-y",
            "-i", temp_in_path,
            "-ar", "16000",      # 16kHz sampling rate
            "-ac", "1",          # Mono channel
            temp_out_path
        ]
        subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
        return temp_out_path
    except Exception as e:
        print("FFmpeg local ASR conversion failed:")
        traceback.print_exc()
        raise e
    finally:
        try:
            if os.path.exists(temp_in_path):
                os.remove(temp_in_path)
        except Exception:
            pass

async def start_tunnel(port: int = 8001) -> tuple[subprocess.Popen, str]:
    """
    Starts a temporary SSH tunnel using Pinggy (primary) or Serveo (fallback).
    Returns the process object and the public HTTPS URL.
    """
    # Try Pinggy first
    cmd = [
        "ssh", "-o", "StrictHostKeyChecking=no",
        "-p", "443", "-R", f"0:127.0.0.1:{port}", "ap@a.pinggy.io"
    ]
    print(f"Starting tunnel with command: {' '.join(cmd)}")
    
    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        stdin=subprocess.DEVNULL,
        text=True,
        bufsize=1
    )
    
    public_url = None
    start_time = time.time()
    
    while time.time() - start_time < 8:
        await asyncio.sleep(0.1)
        if proc.poll() is not None:
            break
            
        line = proc.stdout.readline()
        if not line:
            continue
        print(f"Tunnel stdout: {line.strip()}")
        
        # Match pinggy https link
        match = re.search(r"https://[a-zA-Z0-9.-]+\.pinggy-free\.link", line)
        if not match:
            match = re.search(r"https://[a-zA-Z0-9.-]+\.pinggy\.net", line)
            
        if match:
            public_url = match.group(0)
            print(f"Tunnel established successfully: {public_url}")
            return proc, public_url
            
    # Try Serveo as fallback
    print("Pinggy failed or timed out. Trying Serveo...")
    try:
        proc.terminate()
    except Exception:
        pass
        
    cmd = [
        "ssh", "-o", "StrictHostKeyChecking=no",
        "-R", f"80:127.0.0.1:{port}", "serveo.net"
    ]
    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        stdin=subprocess.DEVNULL,
        text=True,
        bufsize=1
    )
    
    start_time = time.time()
    while time.time() - start_time < 8:
        await asyncio.sleep(0.1)
        if proc.poll() is not None:
            break
            
        line = proc.stdout.readline()
        if not line:
            continue
        print(f"Tunnel stdout: {line.strip()}")
        
        match = re.search(r"https://[a-zA-Z0-9.-]+\.serveousercontent\.com", line)
        if not match:
            match = re.search(r"https://[a-zA-Z0-9.-]+\.serveo\.net", line)
            
        if match:
            public_url = match.group(0)
            print(f"Tunnel established successfully (Serveo): {public_url}")
            return proc, public_url
            
    try:
        proc.terminate()
    except Exception:
        pass
    raise Exception("Failed to establish a public tunnel for audio download.")

@app.get("/v1/models")
async def list_models(authorization: str = Header(None)):
    """
    TypeWhisper calls /v1/models to verify the API connection and populate model dropdowns.
    We forward this to your workspace-specific Alibaba Cloud endpoint.
    """
    if authorization:
        try:
            with open("/Users/keira/Documents/GitHub/qwen_proxy/api_key.txt", "w") as f:
                f.write(authorization)
        except Exception:
            pass
    headers = {}
    if authorization:
        headers["Authorization"] = authorization
        
    async with httpx.AsyncClient(verify=False) as client:
        try:
            response = await client.get(
                f"{ALIBABA_BASE_URL}/models",
                headers=headers,
                timeout=10.0
            )
            if response.status_code == 200:
                models_data = response.json()
                if isinstance(models_data, dict) and "data" in models_data and isinstance(models_data["data"], list):
                    allowed_fun_asr = {
                        "fun-asr-mtl",
                        "fun-asr-mtl-spk2",
                        "fun-asr-mtl-spk3",
                        "fun-asr-mtl-spk4",
                        "fun-asr-mtl-spk5",
                        "fun-asr-mtl-spk6"
                    }
                    filtered_data = []
                    for m in models_data["data"]:
                        if not isinstance(m, dict):
                            continue
                        mid = m.get("id")
                        if not mid:
                            continue
                        # Keep only non-realtime: fun-asr-mtl series, qwen3.5-omni, qwen3.6, qwen3.7, and qwen3-asr-flash series
                        is_allowed = (
                            (mid in allowed_fun_asr or
                             "qwen3.5-omni" in mid.lower() or
                             "qwen3.6" in mid.lower() or
                             "qwen3.7" in mid.lower() or
                             "qwen3-asr-flash" in mid.lower()) and
                            "realtime" not in mid.lower()
                        )
                        if is_allowed:
                            filtered_data.append(m)
                            
                    model_ids = {m.get("id") for m in filtered_data if isinstance(m, dict)}
                    custom_models = [
                        {"id": "fun-asr-mtl", "object": "model", "owned_by": "alibaba"},
                        {"id": "fun-asr-mtl-spk2", "object": "model", "owned_by": "alibaba"},
                        {"id": "fun-asr-mtl-spk3", "object": "model", "owned_by": "alibaba"},
                        {"id": "fun-asr-mtl-spk4", "object": "model", "owned_by": "alibaba"},
                        {"id": "fun-asr-mtl-spk5", "object": "model", "owned_by": "alibaba"},
                        {"id": "fun-asr-mtl-spk6", "object": "model", "owned_by": "alibaba"},
                        {"id": "qwen3-asr-flash", "object": "model", "owned_by": "alibaba"},
                        {"id": "qwen3.5-omni-flash", "object": "model", "owned_by": "alibaba"}
                    ]
                    for custom_model in custom_models:
                        if custom_model["id"] not in model_ids:
                            filtered_data.append(custom_model)
                            
                    # Append local models
                    for local_model in LOCAL_MODELS:
                        filtered_data.append(local_model)
                        
                    # Append Gemini models
                    for gemini_model in GEMINI_MODELS:
                        filtered_data.append(gemini_model)
                        
                    # Fetch and append local Ollama models if running
                    try:
                        async with httpx.AsyncClient() as client:
                            ollama_resp = await client.get("http://127.0.0.1:11434/v1/models", timeout=1.5)
                            if ollama_resp.status_code == 200:
                                ollama_models = ollama_resp.json().get("data", [])
                                for om in ollama_models:
                                    if isinstance(om, dict) and om.get("id"):
                                        # Tag owned_by as ollama for clarity
                                        om["owned_by"] = "ollama"
                                        filtered_data.append(om)
                    except Exception:
                        pass
                        
                    models_data["data"] = filtered_data
                return models_data
        except Exception as e:
            print("Failed to fetch models from Alibaba Cloud, using mock fallback:")
            traceback.print_exc()
            
    # Mock fallback list containing Qwen models, Fun-ASR models, and local models
    mock_data = [
        {"id": "qwen3.5-omni-flash", "object": "model", "owned_by": "alibaba"},
        {"id": "qwen3.5-omni-plus", "object": "model", "owned_by": "alibaba"},
        {"id": "fun-asr-mtl", "object": "model", "owned_by": "alibaba"},
        {"id": "fun-asr-mtl-spk2", "object": "model", "owned_by": "alibaba"},
        {"id": "fun-asr-mtl-spk3", "object": "model", "owned_by": "alibaba"},
        {"id": "fun-asr-mtl-spk4", "object": "model", "owned_by": "alibaba"},
        {"id": "fun-asr-mtl-spk5", "object": "model", "owned_by": "alibaba"},
        {"id": "fun-asr-mtl-spk6", "object": "model", "owned_by": "alibaba"},
        {"id": "qwen3-asr-flash", "object": "model", "owned_by": "alibaba"},
        {"id": "qwen3.7-plus", "object": "model", "owned_by": "alibaba"},
        {"id": "qwen3.6-flash", "object": "model", "owned_by": "alibaba"}
    ]
    mock_data.extend(LOCAL_MODELS)
    mock_data.extend(GEMINI_MODELS)
    
    # Fetch and append local Ollama models in mock fallback
    try:
        response = httpx.get("http://127.0.0.1:11434/v1/models", timeout=1.0)
        if response.status_code == 200:
            ollama_models = response.json().get("data", [])
            for om in ollama_models:
                if isinstance(om, dict) and om.get("id"):
                    om["owned_by"] = "ollama"
                    mock_data.append(om)
    except Exception:
        pass
        
    return {
        "object": "list",
        "data": mock_data
    }

@app.post("/v1/audio/transcriptions")
async def transcribe(
    file: UploadFile = File(...),
    model: str = Form(...),
    language: str = Form(None),
    authorization: str = Header(None)
):
    save_last_used_model("asr", model)
    # 1. Read file bytes
    file_bytes = await file.read()
    file_size_mb = len(file_bytes) / (1024 * 1024)
    print(f"==========================================================")
    print(f"Proxy received file: '{file.filename}'")
    print(f"File size: {file_size_mb:.2f} MB")
    print(f"Requested model: '{model}'")
    print(f"==========================================================")
    
    if not file_bytes:
        raise HTTPException(status_code=400, detail="Empty audio file provided.")
        
    filename = file.filename or "audio.wav"
    file_format = filename.split(".")[-1].lower()
    
    # Parse speaker count hint if specified in the model name (e.g. local-zipformer-vi-30m-punc-diar-spk2)
    speaker_count = None
    model_clean = model
    match = re.search(r"[-:]spk(\d+)|[-:](\d+)spk", model.lower())
    if match:
        spk_str = match.group(1) or match.group(2)
        if spk_str:
            speaker_count = int(spk_str)
            model_clean = model[:match.start()]
            print(f"Detected speaker_count={speaker_count} from model name. Cleaned model name: '{model_clean}'.")
            
    # Check if a local offline model is requested
    if model_clean.lower().startswith("local-zipformer-"):
        # Convert input audio to 16kHz WAV for the local engine
        temp_wav_path = convert_to_local_wav(file_bytes, file_format)
        
        enable_punctuation = "-punc" in model_clean.lower()
        enable_diarization = "-diar" in model_clean.lower()
        
        if "30m" in model_clean.lower():
            model_name = "zipformer-30m-rnnt-6000h"
        else:
            model_name = "sherpa-onnx-zipformer-vi-2025-04-20"
            
        base_dir = os.path.dirname(os.path.abspath(__file__))
        model_path = os.path.join(base_dir, "models", model_name)
        
        # Verify model files exist
        if not os.path.exists(model_path):
            try:
                if os.path.exists(temp_wav_path):
                    os.remove(temp_wav_path)
            except Exception:
                pass
            raise HTTPException(
                status_code=400,
                detail=f"Local ASR model '{model_name}' not found at {model_path}. Please run 'python3 download_models.py' to download it."
            )
            
        if enable_punctuation:
            punc_path = os.path.join(base_dir, "models", "vibert-capu")
            if not os.path.exists(punc_path):
                try:
                    if os.path.exists(temp_wav_path):
                        os.remove(temp_wav_path)
                except Exception:
                    pass
                raise HTTPException(
                    status_code=400,
                    detail=f"Punctuation model not found at models/vibert-capu. Please run 'python3 download_models.py' to download it."
                )
                
        if enable_diarization:
            diar_path = os.path.join(base_dir, "models", "pyannote-onnx")
            if not os.path.exists(os.path.join(diar_path, "segmentation-community-1.onnx")) or not os.path.exists(os.path.join(diar_path, "embedding_encoder.onnx")):
                try:
                    if os.path.exists(temp_wav_path):
                        os.remove(temp_wav_path)
                except Exception:
                    pass
                raise HTTPException(
                    status_code=400,
                    detail=f"Local diarization models not found at models/pyannote-onnx. Please run 'python3 download_models.py' to download them."
                )
                
        try:
            from core.asr_engine import TranscriberPipeline
            
            # Set up config
            config = {
                'cpu_threads': 4,
                'restore_punctuation': enable_punctuation,
                'speaker_diarization': enable_diarization,
                'save_ram': True,
                'bypass_vad': False
            }
            
            if enable_diarization:
                config['speaker_model'] = 'community1_onnx'
                if speaker_count is not None:
                    config['num_speakers'] = speaker_count
            
            print(f"Running local ASR pipeline (model={model_name}, punctuation={enable_punctuation}, diarization={enable_diarization})...")
            pipeline = TranscriberPipeline(
                file_path=temp_wav_path,
                model_path=model_path,
                config=config
            )
            
            result_data = pipeline.run()
            
            if enable_diarization:
                sentences = []
                for seg in result_data.get("segments", []):
                    text = seg.get("text", "").strip()
                    speaker_id = seg.get("speaker_id")
                    if speaker_id is not None:
                        sentences.append(f"Speaker {speaker_id}: {text}")
                    else:
                        sentences.append(text)
                final_text = "\n".join(sentences)
            else:
                final_text = result_data["text"]
                
            print("Local ASR transcription completed successfully!")
            return {"text": final_text}
            
        except Exception as e:
            print("Exception during local ASR process:")
            traceback.print_exc()
            if isinstance(e, HTTPException):
                raise e
            raise HTTPException(status_code=500, detail=str(e))
            
        finally:
            # Clean up temp WAV file
            try:
                if os.path.exists(temp_wav_path):
                    os.remove(temp_wav_path)
            except Exception:
                pass
        
    filename = file.filename or "audio.wav"
    file_format = filename.split(".")[-1].lower()
    
    # Compress the audio to ensure it's tiny for transmission
    compressed_bytes, file_format = compress_audio(file_bytes, file_format)
    
    # Extract file extension and format
    if file_format == "m4a":
        mime_type = "audio/x-m4a"
    elif file_format == "mp3":
        mime_type = "audio/mpeg"
    elif file_format in ["wav", "ogg", "aac", "flac"]:
        mime_type = f"audio/{file_format}"
    else:
        mime_type = "audio/wav"
        file_format = "wav"
            
    # Check if we should use the async Fun-ASR path (for diarization)
    is_async_asr = "fun-asr" in model_clean.lower() and "realtime" not in model_clean.lower() and "flash" not in model_clean.lower()
    
    if is_async_asr:
        # Write compressed audio to our hosted directory
        temp_filename = f"audio_{int(time.time())}.{file_format}"
        temp_filepath = os.path.join(HOSTED_DIR, temp_filename)
        with open(temp_filepath, "wb") as f:
            f.write(compressed_bytes)
            
        tunnel_proc = None
        try:
            # Start the public tunnel
            tunnel_proc, public_url = await start_tunnel()
            audio_url = f"{public_url}/files/{temp_filename}"
            print(f"Public audio URL for DashScope: {audio_url}")
            
            # Workspace ID domain
            WORKSPACE_DOMAIN = "https://ws-fnt3psiwu1morhou.ap-southeast-1.maas.aliyuncs.com"
            
            headers = {
                "Authorization": authorization,
                "Content-Type": "application/json",
                "X-DashScope-Async": "enable"
            }
            
            # Build ASR submit payload
            submit_payload = {
                "model": model_clean,
                "input": {
                    "file_urls": [audio_url]
                },
                "parameters": {
                    "diarization_enabled": True
                }
            }
            if speaker_count is not None:
                submit_payload["parameters"]["speaker_count"] = speaker_count
            if language:
                submit_payload["parameters"]["language_hints"] = [language]
                
            print("Submitting async ASR task to Alibaba Cloud...")
            async with httpx.AsyncClient(verify=False) as client:
                submit_resp = await client.post(
                    f"{WORKSPACE_DOMAIN}/api/v1/services/audio/asr/transcription",
                    json=submit_payload,
                    headers=headers,
                    timeout=30.0
                )
                
                print(f"Alibaba submit status code: {submit_resp.status_code}")
                if submit_resp.status_code != 200:
                    print(f"Alibaba ASR submit failed: {submit_resp.text}")
                    raise HTTPException(
                        status_code=submit_resp.status_code,
                        detail=f"Alibaba ASR submit failed: {submit_resp.text}"
                    )
                    
                task_data = submit_resp.json()
                task_id = task_data.get("output", {}).get("task_id")
                if not task_id:
                    raise HTTPException(
                        status_code=500,
                        detail=f"Failed to obtain task_id from Alibaba: {task_data}"
                    )
                    
                print(f"ASR Task submitted successfully. Task ID: {task_id}")
                
                # Poll task status
                finished = False
                transcription_url = None
                poll_start_time = time.time()
                
                # Allow polling up to 25 minutes
                while not finished and (time.time() - poll_start_time < 1500):
                    await asyncio.sleep(3.0)
                    
                    query_resp = await client.get(
                        f"{WORKSPACE_DOMAIN}/api/v1/tasks/{task_id}",
                        headers={"Authorization": authorization},
                        timeout=20.0
                    )
                    
                    if query_resp.status_code != 200:
                        print(f"Failed to query task status: {query_resp.text}")
                        continue
                        
                    query_data = query_resp.json()
                    status = query_data.get("output", {}).get("task_status", "UNKNOWN")
                    print(f"Polling task {task_id}... Status: {status}")
                    
                    if status == "SUCCEEDED":
                        results = query_data.get("output", {}).get("results", [])
                        if results and results[0].get("subtask_status") == "SUCCEEDED":
                            transcription_url = results[0].get("transcription_url")
                            finished = True
                        else:
                            subtask_err = results[0] if results else {}
                            raise HTTPException(
                                status_code=500,
                                detail=f"ASR subtask failed: {subtask_err.get('message', 'Unknown error')} (Code: {subtask_err.get('code')})"
                            )
                    elif status in ["FAILED", "UNKNOWN"]:
                        raise HTTPException(
                            status_code=500,
                            detail=f"ASR task failed with status: {status}. Full response: {query_data}"
                        )
                        
                if not finished:
                    raise HTTPException(
                        status_code=504,
                        detail="ASR task timed out on Alibaba Cloud."
                    )
                    
                # Download results
                print("ASR task succeeded! Downloading transcription results...")
                result_resp = await client.get(transcription_url, timeout=30.0)
                if result_resp.status_code != 200:
                    raise HTTPException(
                        status_code=502,
                        detail=f"Failed to download transcription results: {result_resp.text}"
                    )
                    
                result_json = result_resp.json()
                
                # Parse sentences and format with speaker labels
                sentences = []
                transcripts = result_json.get("transcripts", [])
                for transcript in transcripts:
                    for sentence in transcript.get("sentences", []):
                        text = sentence.get("text", "").strip()
                        speaker_id = sentence.get("speaker_id")
                        
                        if speaker_id is not None:
                            sentences.append(f"Speaker {speaker_id}: {text}")
                        else:
                            sentences.append(text)
                            
                final_text = "\n".join(sentences)
                print("Diarized transcription completed successfully!")
                return {"text": final_text}
                
        except Exception as e:
            print("Exception during async ASR process:")
            traceback.print_exc()
            if isinstance(e, HTTPException):
                raise e
            raise HTTPException(status_code=500, detail=str(e))
            
        finally:
            # Terminate tunnel and clean up the temp file
            if tunnel_proc:
                print("Stopping tunnel...")
                try:
                    tunnel_proc.terminate()
                except Exception:
                    pass
            if os.path.exists(temp_filepath):
                print("Deleting temporary audio file...")
                try:
                    os.remove(temp_filepath)
                except Exception:
                    pass
                    
    # Otherwise, run standard OpenAI-compatible Direct base64 path (Qwen Omni / Qwen ASR Flash)
    base64_audio = base64.b64encode(compressed_bytes).decode("utf-8")
    
    content_items = [
        {
            "type": "input_audio",
            "input_audio": {
                "data": f"data:{mime_type};base64,{base64_audio}",
                "format": file_format
            }
        }
    ]
    
    # Only add the text instruction prompt for Omni (multimodal/chat) models.
    # Dedicated ASR models (like qwen3-asr-flash) do not support text prompts.
    is_asr_model = "asr" in model_clean.lower()
    if not is_asr_model:
        content_items.append({
            "type": "text",
            "text": "Please transcribe this audio exactly in its original language, without translating. Do not add any introductory or concluding remarks."
        })

    payload = {
        "model": model_clean,
        "messages": [
            {
                "role": "user",
                "content": content_items
            }
        ]
    }
    
    asr_options = {"enable_itn": False}
    if language:
        asr_options["language"] = language
        
    payload["extra_body"] = {"asr_options": asr_options}
    
    headers = {
        "Content-Type": "application/json"
    }
    if authorization:
        headers["Authorization"] = authorization
        
    async with httpx.AsyncClient(verify=False) as client:
        try:
            print("Sending base64 request to Alibaba Cloud API...")
            response = await client.post(
                f"{ALIBABA_BASE_URL}/chat/completions",
                json=payload,
                headers=headers,
                timeout=120.0
            )
            
            print(f"Alibaba Cloud API response status: {response.status_code}")
            if response.status_code != 200:
                print(f"Alibaba Cloud API error response: {response.text}")
                raise HTTPException(
                    status_code=response.status_code, 
                    detail=f"Alibaba API error: {response.text}"
                )
                
            result = response.json()
            print(f"Alibaba response JSON: {result}")
            transcription_text = result["choices"][0]["message"]["content"]
            print("Transcription completed successfully!")
            return {"text": transcription_text}
            
        except httpx.HTTPError as e:
            print("Failed to connect to Alibaba Cloud API (HTTPError):")
            traceback.print_exc()
            raise HTTPException(
                status_code=502, 
                detail=f"Failed to connect to Alibaba Cloud API: {str(e)} ({type(e).__name__})"
            )
        except KeyError:
            print("Unexpected response format from Alibaba Cloud:")
            traceback.print_exc()
            raise HTTPException(status_code=502, detail=f"Unexpected response format from Alibaba: {result}")
        except Exception as e:
            print("General Proxy Exception:")
            traceback.print_exc()
            raise HTTPException(status_code=500, detail=str(e))

@app.post("/v1/chat/completions")
async def chat_completions(
    request: Request,
    authorization: str = Header(None)
):
    """
    Proxy endpoint for LLM workflows. Intercepts the request to force a high max_tokens limit 
    (avoiding prompt truncations in TypeWhisper) and forwards it to Gemini's official OpenAI endpoint.
    """
    try:
        body = await request.json()
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid JSON body")
        
    print("----------------------------------------------------------")
    print("Intercepted LLM chat completion request:")
    print(f"Original model: {body.get('model')}")
    print("----------------------------------------------------------")
    
    # Force max_tokens and max_completion_tokens to a high value so the response is not truncated
    if "max_tokens" in body:
        print(f"Overriding max_tokens from {body['max_tokens']} -> 8192")
        body["max_tokens"] = 8192
    if "max_completion_tokens" in body:
        print(f"Overriding max_completion_tokens from {body['max_completion_tokens']} -> 8192")
        body["max_completion_tokens"] = 8192
        
    model = body.get("model", "")
    save_last_used_model("llm", model)
    
    headers = {
        "Content-Type": "application/json"
    }
    
    # Router to the correct API provider
    if "gemini" in model.lower():
        # Forward to Google Gemini's official OpenAI-compatible endpoint
        url = "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions"
        if authorization:
            api_key = authorization.replace("Bearer ", "").strip()
            headers["Authorization"] = f"Bearer {api_key}"
    else:
        # Check if the requested model is running on local Ollama
        is_ollama = False
        try:
            response = httpx.get("http://127.0.0.1:11434/v1/models", timeout=1.0)
            if response.status_code == 200:
                ollama_models = [m.get("id") for m in response.json().get("data", []) if isinstance(m, dict)]
                if model in ollama_models:
                    is_ollama = True
        except Exception:
            pass
            
        if is_ollama:
            url = "http://127.0.0.1:11434/v1/chat/completions"
            # Note: local Ollama does not require authorization
        else:
            # Otherwise, route to Alibaba Cloud compatible endpoint
            url = "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"
            if authorization:
                headers["Authorization"] = authorization

    print(f"Forwarding LLM completions request to: {url}")
    
    async with httpx.AsyncClient(verify=False) as client:
        try:
            response = await client.post(
                url,
                json=body,
                headers=headers,
                timeout=120.0
            )
            print(f"Provider API response status: {response.status_code}")
            return Response(
                content=response.content,
                status_code=response.status_code,
                headers={
                    "Content-Type": response.headers.get("Content-Type", "application/json"),
                    "Transfer-Encoding": response.headers.get("Transfer-Encoding", "")
                }
            )
        except httpx.HTTPError as e:
            print("Failed to forward chat completion request (HTTPError):")
            traceback.print_exc()
            raise HTTPException(status_code=502, detail=f"API provider connection failed: {str(e)}")
def _check_for_updates_async():
    # Wait a few seconds after startup before checking to avoid slowing down boot
    time.sleep(5)
    try:
        import subprocess
        # Get the remote origin URL dynamically
        res = subprocess.run(["git", "config", "--get", "remote.origin.url"], capture_output=True, text=True)
        url = res.stdout.strip()
        if not url:
            return
            
        # Parse owner and repo name
        match = re.search(r"github\.com[:/]([^/]+)/([^.]+)", url)
        if not match:
            return
            
        owner, repo = match.groups()
        raw_url = f"https://raw.githubusercontent.com/{owner}/{repo}/main/version.json"
        
        req = urllib.request.Request(raw_url)
        req.add_header("User-Agent", "Mozilla/5.0")
        
        with urllib.request.urlopen(req, timeout=5) as response:
            remote_config = json.loads(response.read().decode('utf-8'))
            remote_version = remote_config.get("version", "")
            
        # Load local version
        local_version_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "version.json")
        if os.path.exists(local_version_path):
            with open(local_version_path, "r") as f:
                local_config = json.load(f)
                local_version = local_config.get("version", "")
                
            if local_version and remote_version and remote_version != local_version:
                try:
                    local_parts = [int(p) for p in local_version.split(".")]
                    remote_parts = [int(p) for p in remote_version.split(".")]
                    if remote_parts > local_parts:
                        print(f"[Update Check] New version available: {remote_version} (local: {local_version})")
                        # Escape quotes for AppleScript
                        msg = f"Có bản cập nhật mới (v{remote_version}). Hãy chạy file update.command để nâng cấp."
                        os.system(f'osascript -e \'display notification "{msg}" with title "TypeWhisper Update Available"\'')
                except Exception:
                    pass
    except Exception as e:
        print(f"[Update Check] Failed: {e}")

@app.on_event("startup")
def startup_event():
    import threading
    threading.Thread(target=_check_for_updates_async, daemon=True).start()

if __name__ == "__main__":
    import uvicorn
    # Start server on localhost port 8001
    uvicorn.run(app, host="127.0.0.1", port=8001)