import subprocess
import sys
import os

def install(package, extra_args=None):
    cmd = [sys.executable, "-m", "pip", "install", package]
    if extra_args:
        cmd.extend(extra_args)
    subprocess.check_call(cmd)

print("Asennetaan paketit...")

# Gradio on jo asennettu
install("huggingface_hub")

# llama-cpp-python CPU-versiolla (tärkeä korjaus)
install("llama-cpp-python", ["--extra-index-url", "https://abetlen.github.io/llama-cpp-python/whl/cpu"])

from huggingface_hub import hf_hub_download
import gradio as gr
from llama_cpp import Llama

# Malli (Gemma 4 E4B - paras sun 2vCPU + 16GB setupille)
MODEL_REPO = "unsloth/gemma-4-E4B-it-GGUF"
MODEL_FILE = "gemma-4-E4B-it-Q4_K_M.gguf"
MODEL_DIR = "./models"

os.makedirs(MODEL_DIR, exist_ok=True)
model_path = os.path.join(MODEL_DIR, MODEL_FILE)

if not os.path.exists(model_path):
    print("Ladataan malli Hugging Facesta (ensimmäinen kerta voi kestää muutaman minuutin)...")
    hf_hub_download(
        repo_id=MODEL_REPO,
        filename=MODEL_FILE,
        local_dir=MODEL_DIR,
        local_dir_use_symlinks=False
    )
    print("Malli ladattu!")

print("Ladataan malli muistiin...")
llm = Llama(
    model_path=model_path,
    n_ctx=4096,
    n_threads=2,        # sun 2 vCPU
    n_gpu_layers=0,     # CPU only
    verbose=False
)

def chat(message, history):
    output = llm.create_chat_completion(
        messages=[{"role": "user", "content": message}],
        max_tokens=512,
        temperature=0.7,
    )
    return output["choices"][0]["message"]["content"]

print("Käynnistetään Gradio...")
gr.ChatInterface(
    chat,
    title="Gemma 4 E4B - Paikallinen chat",
    description="Toimii sun laitteella. Puhut suoraan mallin kanssa selaimessa."
).launch(server_name="0.0.0.0", server_port=7860)