import os
import subprocess
import gradio as gr
import spaces

MODEL_REPO = "HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive"
MODEL_FILE = "Gemma-4-E4B-Uncensored-HauhauCS-Aggressive-Q5_K_P.gguf"
MMPROJ_FILE = "mmproj-Gemma-4-E4B-Uncensored-HauhauCS-Aggressive-f16.gguf"

MODEL_PATH = os.path.join("/tmp", MODEL_FILE)
MMPROJ_PATH = os.path.join("/tmp", MMPROJ_FILE)

server_process = None

LLAMA_DIR = "/tmp/llama"
LLAMA_SERVER = os.path.join(LLAMA_DIR, "llama-server")
LLAMA_RELEASE = "b10182"
LLAMA_ARCHIVE = f"llama.cpp-{LLAMA_RELEASE}-cuda-12.8-amd64.tar.gz"
LLAMA_URL = f"https://github.com/ai-dock/llama.cpp-cuda/releases/download/{LLAMA_RELEASE}/{LLAMA_ARCHIVE}"


def download_llama():
    if os.path.exists(LLAMA_SERVER):
        return

    os.makedirs(LLAMA_DIR, exist_ok=True)

    archive = os.path.join("/tmp", LLAMA_ARCHIVE)

    subprocess.run(
        ["wget", "-q", "--show-progress", "-O", archive, LLAMA_URL],
        check=True,
    )

    subprocess.run(
        ["tar", "-xzf", archive, "-C", LLAMA_DIR, "--strip-components=1"],
        check=True,
    )

    if not os.path.exists(LLAMA_SERVER):
        raise RuntimeError(
            f"llama-server was not found after extraction: {LLAMA_SERVER}"
        )

    os.chmod(LLAMA_SERVER, 0o755)


def download_models():
    from huggingface_hub import hf_hub_download

    if not os.path.exists(MODEL_PATH):
        hf_hub_download(
            repo_id=MODEL_REPO,
            filename=MODEL_FILE,
            local_dir="/tmp",
        )

    if not os.path.exists(MMPROJ_PATH):
        hf_hub_download(
            repo_id=MODEL_REPO,
            filename=MMPROJ_FILE,
            local_dir="/tmp",
        )


@spaces.GPU
def generate(prompt, max_tokens, temperature):
    global server_process

    download_llama()
    download_models()

    if server_process is None or server_process.poll() is not None:
        cmd = [
            LLAMA_SERVER,
            "-m", MODEL_PATH,
            "--mmproj", MMPROJ_PATH,
            "--host", "127.0.0.1",
            "--port", "8080",
            "-c", "32768",
            "-n", "4096",
            "--n-gpu-layers", "99",
            "-t", "2",
        ]

        server_env = os.environ.copy()
        server_env["LD_LIBRARY_PATH"] = (
            LLAMA_DIR
            + ":"
            + server_env.get("LD_LIBRARY_PATH", "")
        )

        server_process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            env=server_env,
        )

    import time
    import requests

    for _ in range(120):
        try:
            r = requests.get("http://127.0.0.1:8080/health", timeout=1)
            if r.ok:
                break
        except Exception:
            pass
        time.sleep(1)
    else:
        raise RuntimeError("llama-server did not become ready.")

    response = requests.post(
        "http://127.0.0.1:8080/completion",
        json={
            "prompt": prompt,
            "n_predict": int(max_tokens),
            "temperature": float(temperature),
            "stream": False,
        },
        timeout=600,
    )

    response.raise_for_status()
    return response.json().get("content", "")


with gr.Blocks(title="Gemma-4-E4B Uncensored") as demo:
    gr.Markdown("# Gemma-4-E4B-Uncensored-HauhauCS-Aggressive")

    prompt = gr.Textbox(
        label="Prompt",
        placeholder="Enter your prompt...",
        lines=8,
    )

    with gr.Row():
        max_tokens = gr.Slider(
            minimum=64,
            maximum=4096,
            value=1024,
            step=64,
            label="Max tokens",
        )

        temperature = gr.Slider(
            minimum=0.0,
            maximum=2.0,
            value=0.7,
            step=0.05,
            label="Temperature",
        )

    output = gr.Textbox(
        label="Response",
        lines=20,
    )

    generate_button = gr.Button("Generate", variant="primary")

    generate_button.click(
        generate,
        inputs=[prompt, max_tokens, temperature],
        outputs=output,
    )


if __name__ == "__main__":
    demo.launch(ssr_mode=False)