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

# 1. Config - Verified Filename
REPO_ID = "erdemozkan/YOLO-7B-Qwen-Coder"
GGUF_FILE = "YOLO-7B-Qwen-q4.gguf" 

print(f"🚀 YOLO CODER: Pulling {GGUF_FILE}...")

try:
    model_path = hf_hub_download(repo_id=REPO_ID, filename=GGUF_FILE)
    print(f"✅ Model found: {model_path}")
except Exception as e:
    print(f"❌ Download failed: {e}")
    model_path = None

# 2. Engine Initialization
if model_path:
    llm = Llama(
        model_path=model_path,
        n_ctx=1024,
        n_threads=2, # Optimized for HF Free Tier
        verbose=False
    )
else:
    llm = None

def yoco_heal(broken_code):
    if not llm: return "❌ Engine Offline. Model file not found."
    if not broken_code.strip(): return "⚠️ Paste code."

    prompt = f"<|im_start|>system\nYou are YOLO CODER. Fix the code perfectly.<|im_end|>\n<|im_start|>user\n{broken_code}<|im_end|>\n<|im_start|>assistant\n"
    
    response = llm(prompt, max_tokens=1024, stop=["<|im_end|>"], echo=False)
    return response["choices"][0]["text"].strip()

# 3. UI logic
css = ".gradio-container {background:#000; color:#ffff00; font-family:monospace;} #header {text-align:center; text-shadow: 0 0 10px #ffff00;}"
with gr.Blocks(css=css) as demo:
    gr.Markdown("# YOLO CODER [DOCKER CPU]", elem_id="header")
    with gr.Row():
        in_code = gr.Code(label="INPUT", language="python", lines=12)
        out_code = gr.Code(label="HEALED", language="python", lines=12)
    btn = gr.Button("YOLO IT!")
    btn.click(fn=yoco_heal, inputs=in_code, outputs=out_code)

if __name__ == "__main__":
    # REQUIRED: Port 7860 for Docker on HF
    demo.launch(server_name="0.0.0.0", server_port=7860)