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

# ============================================================
# Config
# ============================================================
MODEL_ID = "BoyBarley/BoyBarley_X"
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,
    device_map="cpu",
)
model.eval()
print("Model loaded successfully.")

# ============================================================
# Chat function (Gradio 5 messages format)
# ============================================================
def respond(message, history, system_prompt, max_tokens):
    if not message or not message.strip():
        return "Please enter a message."
    
    # Build messages list in OpenAI format
    messages = [{"role": "system", "content": system_prompt}]
    
    # history is list of {"role": ..., "content": ...} in Gradio 5
    for msg in history:
        if isinstance(msg, dict):
            messages.append(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
# ============================================================
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-20 seconds. Please be patient! 🙏
"""

demo = gr.ChatInterface(
    fn=respond,
    type="messages",
    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?", DEFAULT_SYSTEM, 150],
        ["List all docker containers", DEFAULT_SYSTEM, 150],
        ["Cek disk usage di /var/log", DEFAULT_SYSTEM, 150],
        ["How to scale a Kubernetes deployment?", DEFAULT_SYSTEM, 150],
    ],
    theme=gr.themes.Soft(),
    cache_examples=False,
)

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