import os, subprocess, urllib.request, tarfile, sys
from huggingface_hub import hf_hub_download

repo = "HauhauCS/Gemma-4-E2B-Uncensored-HauhauCS-Aggressive"
model_file = "Gemma-4-E2B-Uncensored-HauhauCS-Aggressive-Q5_K_P.gguf"
mmproj_file = "mmproj-Gemma-4-E2B-Uncensored-HauhauCS-Aggressive-f16.gguf"
LLAMA_RELEASE = "https://github.com/ggml-org/llama.cpp/releases/download/b9294/llama-b9294-bin-ubuntu-x64.tar.gz"

BIN_DIR = "llama-bin"
PORT = 7860

if not os.path.exists(BIN_DIR):
    os.makedirs(BIN_DIR)
    tar_path = os.path.join(BIN_DIR, "llama.tar.gz")
    print("Downloading llama-server binary...")
    urllib.request.urlretrieve(LLAMA_RELEASE, tar_path)
    with tarfile.open(tar_path, "r:gz") as tar:
        tar.extractall(path=BIN_DIR)
    os.remove(tar_path)

llama_dir = [d for d in os.listdir(BIN_DIR) if d.startswith("llama-")][0]
llama_server = os.path.join(BIN_DIR, llama_dir, "llama-server")
os.chmod(llama_server, 0o755)

print("Downloading model...")
model_path = hf_hub_download(repo, model_file)
mmproj_path = hf_hub_download(repo, mmproj_file)

env = os.environ.copy()
env["LD_LIBRARY_PATH"] = os.path.join(BIN_DIR, llama_dir) + ":" + env.get("LD_LIBRARY_PATH", "")

cmd = [
    llama_server,
    "-m", model_path,
    "--mmproj", mmproj_path,
    "--host", "0.0.0.0",
    "--port", str(PORT),
    "-c", "4096",
    "-t", "2",
    "--no-mmap",
    "--reasoning", "off"
]

print(f"Starting llama-server on port {PORT}...")
subprocess.run(cmd, env=env)