#!/usr/bin/env python3
"""Optional Gradio UI for a PRO/paid CPU Space. Static Space uses index.html."""
import os, subprocess, tempfile
import gradio as gr
from huggingface_hub import hf_hub_download

MODEL_REPO = os.environ.get('MODEL_REPO', 'zerolimites/BitNet-TinyDialogues-1.58b')
GGUF_FILE = os.environ.get('GGUF_FILE', 'gguf/bitnet-tinydialogues-TQ1_0.gguf')
LLAMA_CLI = os.environ.get('LLAMA_CLI', 'llama-cli')

def load_model():
    return hf_hub_download(MODEL_REPO, GGUF_FILE)

def chat(message, history):
    model = load_model()
    prompt = 'User: ' + message + ' Assistant:'
    cmd = [LLAMA_CLI, '-m', model, '-p', prompt, '-n', '64', '-c', '128',
           '--temp', '0.8', '-t', '2', '-r', 'User:', '--no-display-prompt', '--simple-io']
    out = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
    return out.stdout.strip() or out.stderr.strip()

demo = gr.ChatInterface(fn=chat, title='BitNet TinyDialogues CPU')
if __name__ == '__main__':
    demo.launch(server_name='0.0.0.0', server_port=7860)