import os, asyncio, uuid, time, uvicorn, random, shutil, tempfile, httpx, re, json
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request, UploadFile, File, BackgroundTasks
from fastapi.responses import JSONResponse, HTMLResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware
from faster_whisper import WhisperModel
from gtts import gTTS

# --- KONFIGURASI CACHE PERSISTEN ---
CACHE_DIR = "/home/user/app/models_cache"
os.makedirs(CACHE_DIR, exist_ok=True)
os.environ["HF_HOME"] = os.path.join(CACHE_DIR, "huggingface")
os.environ["OLLAMA_MODELS"] = os.path.join(CACHE_DIR, ".ollama")

# --- STRATEGI KONFIGURASI MODEL ---
MODEL_STRATEGY = {
    "deepseek": [
        {"provider": "ollama", "model": "deepseek-r1:7b"},
        {"provider": "openrouter", "model": "deepseek/deepseek-r1"}
    ],
    "qwen": [
        {"provider": "ollama", "model": "qwen2.5:7b-instruct"},
        {"provider": "openrouter", "model": "qwen/qwen-2.5-7b-instruct"}
    ],
    "gpt": [{"provider": "openrouter", "model": "openai/gpt-4o-mini"}],
    "gemini": [{"provider": "openrouter", "model": "google/gemini-2.0-flash"}]
}
QUEUE = asyncio.Queue()

# Load Whisper Base secara global
print(f"[*] Memuat model Whisper BASE dari {CACHE_DIR}...")
whisper_model = WhisperModel("base", device="cpu", compute_type="int8", download_root=CACHE_DIR)

def remove_file(path: str):
    if os.path.exists(path):
        os.remove(path)

def clean_json_output(text: str):
    """Pembersih JSON yang lebih agresif untuk menangani noise dari model reasoning"""
    if not text: return "[]"
    text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
    text = re.sub(r'```json\s*', '', text)
    text = re.sub(r'```\s*', '', text)
    
    match = re.search(r'(\[.*\]|\{.*\})', text, re.DOTALL)
    if not match: return "[]"
    
    candidate = match.group(1).strip()
    try:
        data = json.loads(candidate)
        return json.dumps([data] if isinstance(data, dict) else data)
    except Exception:
        valid_match = re.search(r'(\[.*?\]|\{.*?\})', candidate, re.DOTALL)
        return clean_json_output(valid_match.group(1)) if valid_match else "[]"

# --- PROVIDERS DENGAN RETRY LOGIC ---
async def call_ollama(messages, model):
    print(f"[Ollama] Menghubungi model {model}...")
    for attempt in range(2):
        try:
            async with httpx.AsyncClient(timeout=httpx.Timeout(240.0, connect=10.0)) as client:
                res = await client.post("http://127.0.0.1:11434/api/chat", 
                                        json={"model": model, "messages": messages, "stream": False})
                if res.status_code == 200:
                    print(f"[Ollama] Berhasil mendapatkan respon dari {model}.")
                    return res.json()["message"]["content"]
        except Exception as e:
            print(f"[Ollama] Gagal (Percobaan {attempt+1}): {str(e)}")
            await asyncio.sleep(2)
    return None

async def call_openrouter(messages, model):
    print(f"[OpenRouter] Menghubungi model {model}...")
    api_key = os.getenv("OPENROUTER_KEYS")
    if not api_key: 
        print("[OpenRouter] API KEY tidak ditemukan!")
        return None
    keys = api_key.split(",")
    
    for i in range(3):
        try:
            async with httpx.AsyncClient(timeout=60.0) as client:
                res = await client.post("https://openrouter.ai/api/v1/chat/completions",
                    headers={
                        "Authorization": f"Bearer {random.choice(keys)}", 
                        "HTTP-Referer": "https://huggingface.co/",
                        "Content-Type": "application/json"
                    },
                    json={"model": model, "messages": messages})
                if res.status_code == 200:
                    print(f"[OpenRouter] Berhasil mendapatkan respon dari {model}.")
                    return res.json()["choices"][0]["message"]["content"]
                elif res.status_code == 429:
                    print(f"[OpenRouter] Rate limited, menunggu {2**i} detik...")
                    await asyncio.sleep(2 ** i)
        except Exception as e:
            print(f"[OpenRouter] Gagal: {str(e)}")
            await asyncio.sleep(1)
    return None

async def smart_router(messages, alias, is_json=False):
    strategies = MODEL_STRATEGY.get(alias, MODEL_STRATEGY["deepseek"])
    for attempt in strategies:
        prov, mod = attempt["provider"], attempt["model"]
        raw = await call_ollama(messages, mod) if prov == "ollama" else await call_openrouter(messages, mod)
        if raw:
            return clean_json_output(raw) if is_json else raw
    return "[]" if is_json else "Error: Semua provider (Ollama & OpenRouter) gagal merespon."

async def worker():
    print("[Worker] Antrean AI dimulai...")
    while True:
        messages, alias, is_json, future = await QUEUE.get()
        try:
            print(f"[Worker] Memproses permintaan untuk alias: {alias}")
            result = await smart_router(messages, alias, is_json)
            future.set_result(result)
        except Exception as e:
            print(f"[Worker] Error fatal: {str(e)}")
            future.set_exception(e)
        finally: 
            QUEUE.task_done()

@asynccontextmanager
async def lifespan(app: FastAPI):
    worker_task = asyncio.create_task(worker())
    yield
    worker_task.cancel()

app = FastAPI(title="AI Debug Gateway", lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

# --- API ENDPOINTS ---

@app.get("/v1/models")
async def list_models():
    return {"object": "list", "data": [{"id": "qwen"}, {"id": "deepseek"}, {"id": "gpt"}, {"id": "gemini"}]}

@app.post("/v1/chat/completions")
async def chat(request: Request):
    body = await request.json()
    model_input = body.get("model", "qwen").lower()
    messages = body.get("messages", [])
    
    prompt_str = str(messages).lower()
    is_json = any(x in prompt_str for x in ["json", "highlight", "clipper", "short", "segment"])

    if is_json:
        messages.insert(0, {
            "role": "system", 
            "content": "You are a professional Video Editor AI. Respond ONLY with a JSON LIST of objects: [{'start': 'MM:SS', 'end': 'MM:SS', 'description': '...', 'reason': '...'}]. No conversational text."
        })

    alias = "qwen" if "qwen" in model_input else "deepseek"
    if "gpt" in model_input: alias = "gpt"
    elif "gemini" in model_input: alias = "gemini"

    future = asyncio.get_event_loop().create_future()
    await QUEUE.put((messages, alias, is_json, future))
    
    # Tunggu hasil dari worker
    try:
        result = await asyncio.wait_for(future, timeout=300.0) # Max 5 menit
    except asyncio.TimeoutError:
        result = "Error: Timeout 5 menit tercapai."

    return {
        "id": f"chatcmpl-{uuid.uuid4()}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model_input,
        "choices": [{"message": {"role": "assistant", "content": result}, "finish_reason": "stop"}]
    }

@app.post("/v1/audio/transcriptions")
async def transcribe(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
    tmp_path = os.path.join(tempfile.gettempdir(), f"{uuid.uuid4()}_{file.filename}")
    with open(tmp_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
    
    try:
        print(f"[*] Transkripsi dimulai: {file.filename}")
        segments, info = whisper_model.transcribe(tmp_path, beam_size=5, language="id", vad_filter=True)
        
        segments_list = []
        full_text = ""
        for s in segments:
            segment_dict = {"start": round(s.start, 2), "end": round(s.end, 2), "text": s.text.strip()}
            segments_list.append(segment_dict)
            full_text += s.text.strip() + " "

        background_tasks.add_task(remove_file, tmp_path)
        return {
            "text": full_text.strip() or "[Tidak ada suara]",
            "segments": segments_list,
            "language": info.language,
            "duration": info.duration
        }
    except Exception as e:
        remove_file(tmp_path)
        return JSONResponse({"error": str(e)}, status_code=500)

@app.post("/v1/audio/speech")
async def tts(request: Request, background_tasks: BackgroundTasks):
    body = await request.json()
    text = body.get("input", "")
    tmp_path = os.path.join(tempfile.gettempdir(), f"{uuid.uuid4()}.mp3")
    try:
        gTTS(text=text, lang="id").save(tmp_path)
        background_tasks.add_task(remove_file, tmp_path)
        return FileResponse(tmp_path, media_type="audio/mpeg")
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)

@app.get("/", response_class=HTMLResponse)
async def dashboard():
    return """
    <!DOCTYPE html>
    <html lang="id">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>IMBECOME DEBUGGER</title>
        <script src="https://cdn.tailwindcss.com"></script>
        <style>
            body { background: #0f172a; color: #e2e8f0; font-family: ui-sans-serif, system-ui; }
            .card { background: #1e293b; border: 1px solid #334155; border-radius: 12px; }
            #output { scroll-behavior: smooth; }
            .status-pulse { animation: pulse 2s cubic-bezier(0.4, 0, 0.6, 1) infinite; }
            @keyframes pulse { 0%, 100% { opacity: 1; } 50% { opacity: .5; } }
        </style>
    </head>
    <body class="p-4 md:p-8">
        <div class="max-w-6xl mx-auto">
            <header class="flex flex-col md:flex-row justify-between items-center mb-8 gap-4 border-b border-slate-800 pb-6">
                <div>
                    <h1 class="text-3xl font-black text-sky-400">IMBECOME <span class="text-white">DEBUGGER</span></h1>
                    <p class="text-xs text-slate-500 mt-1">Gateway for YT Short Clipper</p>
                </div>
                <div class="flex items-center gap-3">
                    <div id="loadingIndicator" class="hidden status-pulse px-3 py-1 bg-amber-500/10 text-amber-400 text-[10px] font-bold rounded-full border border-amber-500/20">PROCESSING...</div>
                    <div class="px-3 py-1 bg-emerald-500/10 text-emerald-400 text-xs font-bold rounded-full border border-emerald-500/20">READY</div>
                </div>
            </header>

            <div class="grid grid-cols-1 lg:grid-cols-2 gap-6">
                <div class="card p-6">
                    <h2 class="font-bold mb-4 flex items-center gap-2">LLM Tester</h2>
                    <select id="modelSelect" class="w-full bg-slate-900 border border-slate-700 text-sm p-2 rounded mb-4 outline-none">
                        <option value="qwen">Qwen 2.5 (7B)</option>
                        <option value="deepseek">DeepSeek R1 (7B)</option>
                        <option value="gemini">Gemini Flash</option>
                        <option value="gpt">GPT-4o Mini</option>
                    </select>
                    <textarea id="chatInput" rows="5" class="w-full bg-slate-900 border border-slate-700 rounded-lg p-3 text-sm mb-4 outline-none" placeholder="Masukkan teks..."></textarea>
                    <div class="grid grid-cols-2 gap-3">
                        <button onclick="runTest('chat')" class="bg-slate-700 hover:bg-slate-600 py-2 rounded-lg text-xs font-bold transition">Chat</button>
                        <button onclick="runTest('json')" class="bg-sky-600 hover:bg-sky-500 py-2 rounded-lg text-xs font-bold transition">Extract Highlight</button>
                    </div>
                </div>

                <div class="space-y-6">
                    <div class="card p-6">
                        <h2 class="font-bold mb-4">Whisper STT</h2>
                        <input type="file" id="audioFile" class="text-xs w-full mb-4 text-slate-400">
                        <button onclick="runWhisper()" class="w-full bg-emerald-600 hover:bg-emerald-500 py-2 rounded-lg text-xs font-bold transition">Transcribe</button>
                    </div>
                    <div class="card p-6">
                        <h2 class="font-bold mb-4">gTTS Tester</h2>
                        <div class="flex gap-2">
                            <input id="ttsText" type="text" class="flex-1 bg-slate-900 border border-slate-700 rounded-lg px-3 py-2 text-sm outline-none" value="Cek sound clipper.">
                            <button onclick="runTTS()" class="bg-orange-600 hover:bg-orange-500 px-4 py-2 rounded-lg text-xs font-bold transition">Play</button>
                        </div>
                        <audio id="audioPlayer" controls class="w-full mt-4 hidden h-8"></audio>
                    </div>
                </div>

                <div class="lg:col-span-2 card">
                    <div class="bg-slate-800/50 px-6 py-3 border-b border-slate-700 flex justify-between items-center">
                        <h2 class="text-sm font-bold text-slate-400">Console Log</h2>
                        <button onclick="document.getElementById('output').innerText = ''" class="text-[10px] text-slate-500 hover:text-white uppercase">Clear</button>
                    </div>
                    <pre id="output" class="p-6 h-80 overflow-y-auto text-[11px] text-emerald-400 font-mono leading-relaxed bg-black/20 flex flex-col-reverse"></pre>
                </div>
            </div>
        </div>

        <script>
            const log = (msg, isData = false) => {
                const out = document.getElementById('output');
                const div = document.createElement('div');
                div.className = 'mb-1 border-l-2 border-slate-700 pl-2 py-1';
                const time = new Date().toLocaleTimeString();
                const content = isData ? JSON.stringify(msg, null, 2) : msg;
                div.innerHTML = `<span class="text-slate-500">[${time}]</span> <span class="${isData ? 'text-sky-300' : ''}">${content}</span>`;
                out.prepend(div);
            };

            const toggleLoading = (show) => {
                document.getElementById('loadingIndicator').classList.toggle('hidden', !show);
            };

            async function runTest(type) {
                const model = document.getElementById('modelSelect').value;
                const text = document.getElementById('chatInput').value;
                if(!text) return alert('Input kosong!');

                log(`Memulai ${type} (${model})...`);
                toggleLoading(true);
                try {
                    const res = await fetch('/v1/chat/completions', {
                        method: 'POST',
                        headers: {'Content-Type': 'application/json'},
                        body: JSON.stringify({
                            model: model,
                            messages: [{role: 'user', content: type === 'json' ? "extract highlight: " + text : text}]
                        })
                    });
                    const data = await res.json();
                    log(data, true);
                } catch (err) { log("Error: " + err.message); }
                finally { toggleLoading(false); }
            }

            async function runWhisper() {
                const file = document.getElementById('audioFile').files[0];
                if(!file) return alert('Pilih file!');
                log("Transkripsi dimulai...");
                toggleLoading(true);
                const fd = new FormData();
                fd.append('file', file);
                try {
                    const res = await fetch('/v1/audio/transcriptions', { method: 'POST', body: fd });
                    const data = await res.json();
                    log(data, true);
                } catch (err) { log("Whisper Error: " + err.message); }
                finally { toggleLoading(false); }
            }

            async function runTTS() {
                const text = document.getElementById('ttsText').value;
                log("Memutar suara...");
                const res = await fetch('/v1/audio/speech', {
                    method: 'POST',
                    headers: {'Content-Type': 'application/json'},
                    body: JSON.stringify({ input: text })
                });
                const blob = await res.blob();
                const player = document.getElementById('audioPlayer');
                player.src = URL.createObjectURL(blob);
                player.classList.remove('hidden');
                player.play();
            }
        </script>
    </body>
    </html>
    """

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)