import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# ============================================================
# Config
# ============================================================
MODEL_ID = "BoyBarley/BoyBarley-v32"
DEFAULT_SYSTEM = "Kamu adalah BoyBarley, AI agent DevOps dan automation engineer. Berikan jawaban yang ringkas dan tepat."

# ============================================================
# Load model (once at startup)
# ============================================================
print(f"Loading {MODEL_ID}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32,  # CPU-safe
    device_map="cpu",
)
print("Model loaded.")

# ============================================================
# Chat function
# ============================================================
def respond(message, history, system_prompt, max_tokens):
    if not message or not message.strip():
        return "Please enter a message."
    
    # Build messages
    messages = [{"role": "system", "content": system_prompt}]
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    messages.append({"role": "user", "content": message})
    
    # Apply chat template
    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(prompt, return_tensors="pt")
    
    # Generate
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=int(max_tokens),
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )
    
    response = tokenizer.decode(
        outputs[0][inputs.input_ids.shape[1]:],
        skip_special_tokens=True,
    ).strip()
    
    return response

# ============================================================
# UI — use ChatInterface (simpler, no schema bugs)
# ============================================================
DESCRIPTION = """
# 🤖 BoyBarley — DevOps AI Agent

Bilingual (Indonesian/English) DevOps automation assistant.  
Specialized in CLI, Docker, Kubernetes operations.

**Model:** [BoyBarley-v32](https://huggingface.co/BoyBarley/BoyBarley-v32) (94.4% eval score)

⚠️ Running on CPU — responses take 5-15 seconds. Be patient! 🙏
"""

demo = gr.ChatInterface(
    fn=respond,
    title="BoyBarley DevOps AI",
    description=DESCRIPTION,
    additional_inputs=[
        gr.Textbox(
            value=DEFAULT_SYSTEM,
            label="System Prompt",
            lines=2,
        ),
        gr.Slider(
            minimum=32, maximum=300, value=150, step=16,
            label="Max new tokens",
        ),
    ],
    examples=[
        ["Siapa nama kamu?"],
        ["List all docker containers"],
        ["Cek disk usage di /var/log"],
        ["How to scale a Kubernetes deployment?"],
        ["Buatkan Dockerfile untuk Node.js"],
    ],
    theme=gr.themes.Soft(),
    cache_examples=False,
)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
    )