#!/usr/bin/env python3
"""
GGUF Translate - Translation & chat with any GGUF model via llama.cpp
Features:
- Auto language detection (lingua-py)
- Direct OpenAI API at /v1/* with auto model loading
- Vision support (auto-downloads mmproj for VL models)
- 10x faster downloads via Xet
"""

import gradio as gr
import requests
import json
import subprocess
import os
import time
import threading
from fastapi import Request
from fastapi.responses import StreamingResponse, JSONResponse

# Language detector (lazy load)
_detector = None

def get_detector():
    """Lazy load lingua detector (~15MB model)."""
    global _detector
    if _detector is None:
        from lingua import LanguageDetectorBuilder, Language
        _detector = LanguageDetectorBuilder.from_languages(
            Language.ENGLISH, Language.SPANISH, Language.FRENCH, Language.GERMAN,
            Language.ITALIAN, Language.PORTUGUESE, Language.JAPANESE, Language.KOREAN,
            Language.CHINESE, Language.ARABIC, Language.HINDI, Language.RUSSIAN,
        ).build()
    return _detector

def detect_language(text: str) -> str:
    """Detect language code from text using lingua-py."""
    try:
        detector = get_detector()
        from lingua import Language
        lang = detector.detect_language_of(text)
        if lang is None:
            return "en"
        lang_map = {
            Language.ENGLISH: "en", Language.SPANISH: "es", Language.FRENCH: "fr",
            Language.GERMAN: "de", Language.ITALIAN: "it", Language.PORTUGUESE: "pt",
            Language.JAPANESE: "ja", Language.KOREAN: "ko", Language.CHINESE: "zh",
            Language.ARABIC: "ar", Language.HINDI: "hi", Language.RUSSIAN: "ru",
        }
        return lang_map.get(lang, "en")
    except Exception:
        return "en"

# Current model info
CURRENT_MODEL = {"path": "/models/translategemma-4b-it.Q4_K_M.gguf", "name": "TranslateGemma 4B", "id": "mradermacher/translategemma-4b-it-GGUF:translategemma-4b-it.Q4_K_M.gguf"}
MODEL_LOCK = threading.Lock()

LANGS = {
    "auto": "Auto-detect",
    "en": "English", "es": "Spanish", "fr": "French", "de": "German",
    "it": "Italian", "pt": "Portuguese", "ja": "Japanese", "ko": "Korean",
    "zh": "Chinese", "ar": "Arabic", "hi": "Hindi", "ru": "Russian",
}

EXAMPLE_MODELS = [
    "mradermacher/translategemma-4b-it-GGUF:translategemma-4b-it.Q4_K_M.gguf",
    "Qwen/Qwen3-0.6B-GGUF:Qwen3-0.6B-Q8_0.gguf",
    "bartowski/Llama-3.2-1B-Instruct-GGUF:Llama-3.2-1B-Instruct-Q4_K_M.gguf",
    "bartowski/gemma-2-2b-it-GGUF:gemma-2-2b-it-Q4_K_M.gguf",
    # Vision models (auto-downloads mmproj)
    "Qwen/Qwen3-VL-2B-Instruct-GGUF:Qwen3VL-2B-Instruct-Q4_K_M.gguf",
    "unsloth/Qwen3-VL-4B-Thinking-1M-GGUF:Qwen3-VL-4B-Thinking-1M-UD-Q4_K_XL.gguf",
]


def load_model_sync(model_id: str) -> tuple[bool, str]:
    """Synchronous model loading for API auto-load. Returns (success, message)."""
    global CURRENT_MODEL

    if not model_id or ":" not in model_id:
        return False, "Invalid format. Use: repo_id:filename"

    # Check if already loaded
    if CURRENT_MODEL.get("id") == model_id:
        return True, f"Model already loaded: {model_id}"

    with MODEL_LOCK:
        # Double-check after acquiring lock
        if CURRENT_MODEL.get("id") == model_id:
            return True, f"Model already loaded: {model_id}"

        repo_id, filename = model_id.split(":", 1)
        model_path = f"/models/{filename}"
        mmproj_path = None
        mmproj_filename = None

        # Stop current server
        try:
            subprocess.run(["pkill", "-f", "llama-server"], timeout=10)
            time.sleep(2)
        except Exception:
            pass

        # Check repo for mmproj file
        try:
            from huggingface_hub import list_repo_files, hf_hub_download
            repo_files = list_repo_files(repo_id)
            for f in repo_files:
                if f.lower().startswith("mmproj") and f.endswith(".gguf"):
                    mmproj_filename = f
                    mmproj_path = f"/models/{mmproj_filename}"
                    break
        except Exception:
            from huggingface_hub import hf_hub_download

        # Download main model
        if not os.path.exists(model_path):
            try:
                hf_hub_download(repo_id=repo_id, filename=filename, local_dir="/models")
            except Exception as e:
                return False, f"Download failed: {e}"

        # Download mmproj if found
        if mmproj_filename and mmproj_path and not os.path.exists(mmproj_path):
            try:
                hf_hub_download(repo_id=repo_id, filename=mmproj_filename, local_dir="/models")
            except Exception:
                mmproj_path = None

        # Check if TranslateGemma
        model_lower = filename.lower()
        is_translategemma = "translategemma" in model_lower

        if is_translategemma:
            cmd = [
                "llama-server", "--model", model_path,
                "--host", "127.0.0.1", "--port", "8080",
                "--ctx-size", "4096", "--threads", "2",
                "--no-jinja", "--chat-template", "gemma",
            ]
        else:
            cmd = [
                "llama-server", "-hf", repo_id,
                "--host", "127.0.0.1", "--port", "8080",
                "--ctx-size", "4096", "--threads", "2",
            ]

        if mmproj_path and os.path.exists(mmproj_path):
            cmd.extend(["--mmproj", mmproj_path])

        try:
            subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        except Exception as e:
            return False, f"Failed to start server: {e}"

        # Wait for server
        for i in range(90):
            try:
                r = requests.get("http://127.0.0.1:8080/health", timeout=2)
                if r.status_code == 200:
                    model_type = " (VL)" if mmproj_path else ""
                    CURRENT_MODEL = {"path": model_path, "name": f"{repo_id}/{filename}{model_type}", "id": model_id, "mmproj": mmproj_path}
                    return True, f"Model loaded: {CURRENT_MODEL['name']}"
            except:
                pass
            time.sleep(2)

        return False, "Server started but health check timed out"


def load_model(model_id: str):
    """Load a new GGUF model from HuggingFace. Generator version for UI."""
    global CURRENT_MODEL

    if not model_id or ":" not in model_id:
        yield "❌ Invalid format. Use: repo_id:filename"
        return

    # Check if already loaded
    if CURRENT_MODEL.get("id") == model_id:
        yield f"✅ Model already loaded: {model_id}"
        return

    repo_id, filename = model_id.split(":", 1)
    model_path = f"/models/{filename}"
    mmproj_path = None
    mmproj_filename = None

    yield f"🔄 Loading model: {repo_id}/{filename}"

    yield "⏹️ Stopping current llama-server..."
    try:
        subprocess.run(["pkill", "-f", "llama-server"], timeout=10)
        time.sleep(2)
    except Exception as e:
        yield f"⚠️ Could not stop server: {e}"

    # Check repo for mmproj file
    try:
        from huggingface_hub import list_repo_files, hf_hub_download
        yield f"🔍 Checking repo for vision support (mmproj)..."
        repo_files = list_repo_files(repo_id)
        for f in repo_files:
            if f.lower().startswith("mmproj") and f.endswith(".gguf"):
                mmproj_filename = f
                mmproj_path = f"/models/{mmproj_filename}"
                yield f"👁️ Found vision projector: {mmproj_filename}"
                break
    except Exception as e:
        yield f"⚠️ Could not list repo files: {e}"
        from huggingface_hub import hf_hub_download

    # Download main model
    if not os.path.exists(model_path):
        yield f"📥 Downloading {filename} with Xet (10x faster)..."
        try:
            hf_hub_download(repo_id=repo_id, filename=filename, local_dir="/models")
            yield f"✅ Downloaded {filename}"
        except Exception as e:
            yield f"❌ Download failed: {e}"
            yield "🔄 Restarting with previous model..."
            subprocess.Popen([
                "llama-server", "--model", CURRENT_MODEL["path"],
                "--host", "127.0.0.1", "--port", "8080", "--ctx-size", "4096", "--threads", "2",
                "--no-jinja", "--chat-template", "gemma",
            ])
            return
    else:
        yield f"✅ Model already cached: {filename}"

    # Download mmproj if found
    if mmproj_filename and mmproj_path:
        if not os.path.exists(mmproj_path):
            yield f"📥 Downloading vision projector: {mmproj_filename}..."
            try:
                hf_hub_download(repo_id=repo_id, filename=mmproj_filename, local_dir="/models")
                yield f"✅ Downloaded {mmproj_filename}"
            except Exception as e:
                yield f"⚠️ mmproj download failed: {e} (will try without vision)"
                mmproj_path = None
        else:
            yield f"✅ mmproj already cached: {mmproj_filename}"

    # Check if TranslateGemma
    model_lower = filename.lower()
    is_translategemma = "translategemma" in model_lower

    if is_translategemma:
        cmd = [
            "llama-server", "--model", model_path,
            "--host", "127.0.0.1", "--port", "8080",
            "--ctx-size", "4096", "--threads", "2",
            "--no-jinja", "--chat-template", "gemma",
        ]
    else:
        cmd = [
            "llama-server", "-hf", repo_id,
            "--host", "127.0.0.1", "--port", "8080",
            "--ctx-size", "4096", "--threads", "2",
        ]

    if mmproj_path and os.path.exists(mmproj_path):
        cmd.extend(["--mmproj", mmproj_path])
        yield f"🚀 Starting llama-server with {'gemma template' if is_translategemma else 'jinja'} + vision..."
    else:
        yield f"🚀 Starting llama-server with {'gemma template' if is_translategemma else 'jinja'}..."

    try:
        subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except Exception as e:
        yield f"❌ Failed to start server: {e}"
        return

    yield "⏳ Waiting for model to load..."
    for i in range(90):
        try:
            r = requests.get("http://127.0.0.1:8080/health", timeout=2)
            if r.status_code == 200:
                model_type = " (VL)" if mmproj_path else ""
                CURRENT_MODEL = {"path": model_path, "name": f"{repo_id}/{filename}{model_type}", "id": model_id, "mmproj": mmproj_path}
                yield f"✅ Model loaded: {CURRENT_MODEL['name']}"
                return
        except:
            pass
        time.sleep(2)
        if i % 10 == 9:
            yield f"⏳ Still loading... ({(i+1)*2}s)"

    yield "⚠️ Server started but health check timed out."


def get_current_model():
    return f"**Model:** {CURRENT_MODEL['name']}"


def translate(text: str, source_lang: str = "auto", target_lang: str = "es") -> str:
    """Translate text between languages."""
    if not text.strip():
        return ""

    if source_lang == "auto":
        source_lang = detect_language(text)

    if source_lang == target_lang:
        return text

    src_name = LANGS.get(source_lang, source_lang)
    tgt_name = LANGS.get(target_lang, target_lang)

    try:
        response = requests.post(
            "http://127.0.0.1:8080/v1/chat/completions",
            json={
                "messages": [{"role": "user", "content": f"Translate from {src_name} to {tgt_name}. Output only the translation, nothing else.\n\n{text}"}],
                "max_tokens": 1024,
                "temperature": 0.1,
            },
            timeout=300,
        )
        return response.json().get("choices", [{}])[0].get("message", {}).get("content", "").strip()
    except Exception as e:
        return f"Error: {str(e)}"


def generate_stream(prompt: str, max_tokens: int = 512, temperature: float = 0.7):
    """Generate text with streaming."""
    if not prompt.strip():
        yield ""
        return

    try:
        response = requests.post(
            "http://127.0.0.1:8080/v1/chat/completions",
            json={"messages": [{"role": "user", "content": prompt}], "max_tokens": max_tokens, "temperature": temperature, "stream": True},
            stream=True, timeout=300,
        )
        full = ""
        for line in response.iter_lines():
            if line:
                line = line.decode("utf-8")
                if line.startswith("data: "):
                    data = line[6:]
                    if data == "[DONE]":
                        break
                    try:
                        chunk = json.loads(data)
                        content = chunk.get("choices", [{}])[0].get("delta", {}).get("content", "")
                        if content:
                            full += content
                            yield full
                    except:
                        pass
        if not full:
            yield "[No response]"
    except Exception as e:
        yield f"Error: {str(e)}"


def process_input(text: str, is_translate: bool, src_lang: str, tgt_lang: str, max_tokens: int, temperature: float):
    """Process input based on mode (translate or chat)."""
    if is_translate:
        yield translate(text, src_lang, tgt_lang)
    else:
        for chunk in generate_stream(text, max_tokens, temperature):
            yield chunk


# ============ Gradio UI ============

with gr.Blocks(title="GGUF Translate") as demo:
    gr.Markdown("# 🌐 GGUF Translate\nTranslation & chat with any GGUF model. **Direct OpenAI API:** `/v1/chat/completions`")

    with gr.Tabs():
        with gr.TabItem("🎯 Main"):
            # Model row
            with gr.Row():
                model_dropdown = gr.Dropdown(choices=EXAMPLE_MODELS, label="Model", allow_custom_value=True, value=EXAMPLE_MODELS[0], scale=3)
                load_btn = gr.Button("🔄 Load", variant="secondary", scale=1)
            current_model_display = gr.Markdown(f"**Model:** {CURRENT_MODEL['name']}")
            load_status = gr.Textbox(label="Status", lines=3, interactive=False, visible=False)

            # Mode toggle
            is_translate = gr.Checkbox(label="Translation mode", value=True, info="Uncheck for chat mode")

            # Language selectors (visible only in translate mode)
            with gr.Row() as lang_row:
                src_lang = gr.Dropdown(choices=["auto", "en", "es", "fr", "de", "it", "pt", "ja", "ko", "zh", "ar", "hi", "ru"], value="auto", label="Source", scale=1)
                tgt_lang = gr.Dropdown(choices=["en", "es", "fr", "de", "it", "pt", "ja", "ko", "zh", "ar", "hi", "ru"], value="es", label="Target", scale=1)

            # Generation params (visible only in chat mode)
            with gr.Row(visible=False) as gen_row:
                max_tok = gr.Slider(64, 2048, 512, step=64, label="Max Tokens")
                temp = gr.Slider(0, 2, 0.7, step=0.1, label="Temperature")

            # Input/Output
            input_text = gr.Textbox(label="Input", lines=5, placeholder="Text to translate or chat prompt...")
            output_text = gr.Textbox(label="Output", lines=8, interactive=False)
            submit_btn = gr.Button("Submit", variant="primary")

            # Toggle visibility based on mode
            def toggle_mode(is_trans):
                return gr.update(visible=is_trans), gr.update(visible=not is_trans)

            is_translate.change(fn=toggle_mode, inputs=[is_translate], outputs=[lang_row, gen_row])

            # Load model
            def show_status():
                return gr.update(visible=True)

            def hide_status():
                return gr.update(visible=False)

            load_btn.click(fn=show_status, outputs=[load_status]).then(
                fn=load_model, inputs=[model_dropdown], outputs=[load_status]
            ).then(
                fn=get_current_model, outputs=[current_model_display]
            ).then(
                fn=hide_status, outputs=[load_status]
            )

            # Submit
            submit_btn.click(
                fn=process_input,
                inputs=[input_text, is_translate, src_lang, tgt_lang, max_tok, temp],
                outputs=output_text,
                api_name="process"
            )

        with gr.TabItem("🔌 API"):
            gr.Markdown("""
### Direct OpenAI-compatible API

**Endpoint:** `/v1/chat/completions`

**Auto Model Loading:** Pass `model` as `repo_id:filename` to auto-load!

```bash
# Auto-loads Qwen3-0.6B if not already loaded
curl https://YOUR-SPACE.hf.space/v1/chat/completions \\
  -H "Content-Type: application/json" \\
  -H "Authorization: Bearer YOUR_HF_TOKEN" \\
  -d '{
    "model": "Qwen/Qwen3-0.6B-GGUF:Qwen3-0.6B-Q8_0.gguf",
    "messages": [{"role": "user", "content": "Hello!"}],
    "max_tokens": 100
  }'
```

**Python (OpenAI SDK):**
```python
from openai import OpenAI

client = OpenAI(
    base_url="https://YOUR-SPACE.hf.space/v1",
    api_key="YOUR_HF_TOKEN"
)

# Model field triggers auto-load if different from current
response = client.chat.completions.create(
    model="Qwen/Qwen3-0.6B-GGUF:Qwen3-0.6B-Q8_0.gguf",
    messages=[{"role": "user", "content": "Hello!"}]
)
print(response.choices[0].message.content)
```

**Vision (VL models auto-download mmproj):**
```python
import base64

with open("image.jpg", "rb") as f:
    b64 = base64.b64encode(f.read()).decode()

response = client.chat.completions.create(
    model="Qwen/Qwen3-VL-2B-Instruct-GGUF:Qwen3VL-2B-Instruct-Q4_K_M.gguf",
    messages=[{
        "role": "user",
        "content": [
            {"type": "text", "text": "What's in this image?"},
            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}}
        ]
    }]
)
```

**Translation via API:**
```python
response = client.chat.completions.create(
    model="mradermacher/translategemma-4b-it-GGUF:translategemma-4b-it.Q4_K_M.gguf",
    messages=[{
        "role": "user",
        "content": "Translate from English to Spanish. Output only the translation.\\n\\nHello, how are you?"
    }],
    temperature=0.1
)
```
            """)


# ============ Launch with custom routes ============

if __name__ == "__main__":
    from fastapi import FastAPI
    import uvicorn

    app = FastAPI()

    @app.api_route("/v1/{path:path}", methods=["GET", "POST", "OPTIONS"])
    async def openai_proxy(path: str, request: Request):
        """Proxy /v1/* requests to llama-server with auto model loading."""
        target_url = f"http://127.0.0.1:8080/v1/{path}"
        try:
            body = await request.body() if request.method == "POST" else None
            is_streaming = False
            requested_model = None

            if body:
                try:
                    data = json.loads(body)
                    is_streaming = data.get("stream", False)
                    requested_model = data.get("model", "")
                except:
                    pass

            # Auto-load model if specified in repo_id:filename format
            if requested_model and ":" in requested_model and "/" in requested_model:
                if CURRENT_MODEL.get("id") != requested_model:
                    success, msg = load_model_sync(requested_model)
                    if not success:
                        return JSONResponse(content={"error": f"Failed to load model: {msg}"}, status_code=500)

            if is_streaming:
                def generate():
                    with requests.post(target_url, data=body, headers={"Content-Type": "application/json"}, stream=True, timeout=300) as r:
                        for chunk in r.iter_content(chunk_size=None):
                            if chunk:
                                yield chunk
                return StreamingResponse(generate(), media_type="text/event-stream")
            else:
                if request.method == "POST":
                    resp = requests.post(target_url, data=body, headers={"Content-Type": "application/json"}, timeout=300)
                else:
                    resp = requests.get(target_url, timeout=300)
                return JSONResponse(content=resp.json(), status_code=resp.status_code)
        except Exception as e:
            return JSONResponse(content={"error": str(e)}, status_code=500)

    @app.get("/health")
    async def health_proxy():
        """Proxy health check to llama-server."""
        try:
            resp = requests.get("http://127.0.0.1:8080/health", timeout=5)
            return JSONResponse(content=resp.json(), status_code=resp.status_code)
        except Exception as e:
            return JSONResponse(content={"status": "error", "message": str(e)}, status_code=503)

    # Mount Gradio app
    app = gr.mount_gradio_app(app, demo, path="/")

    uvicorn.run(app, host="0.0.0.0", port=7860)