import gradio as gr
import os
import re
import json

HF_TOKEN = os.environ.get("HF_TOKEN")
if HF_TOKEN:
    from huggingface_hub import login
    login(token=HF_TOKEN, add_to_git_credential=False)

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
import torch

BASE_MODEL = "unsloth/Qwen2.5-Coder-7B-Instruct-bnb-4bit"
LORA_MODEL = "zurd46/Zurd_n8n_Builder"

print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16
    ),
    device_map="auto",
    trust_remote_code=True,
)
model = PeftModel.from_pretrained(model, LORA_MODEL, token=HF_TOKEN)
model.eval()
print("Model loaded!")

def respond(message, history):
    messages = [{"role": "system", "content": "Du bist Zurd n8n Builder. Generiere n8n Workflow JSON in ```json Code-Bloecken."}]
    for h in history:
        messages.append({"role": "user", "content": h[0]})
        if h[1]:
            messages.append({"role": "assistant", "content": h[1]})
    messages.append({"role": "user", "content": message})

    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    outputs = model.generate(
        **inputs,
        max_new_tokens=2048,
        temperature=0.7,
        do_sample=True,
        top_p=0.9,
        pad_token_id=tokenizer.eos_token_id
    )
    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    return response

demo = gr.ChatInterface(
    respond,
    title="Zurd n8n Builder",
    description="KI-Assistent fuer n8n Workflow Automatisierung. Beschreibe deinen gewuenschten Workflow und erhalte n8n JSON.",
    examples=[
        "Erstelle einen Workflow der taeglich eine E-Mail mit dem Wetter sendet",
        "Automatisierung: GitHub Issues -> Slack Benachrichtigung",
        "API Daten abrufen und in Google Sheets speichern"
    ],
    theme=gr.themes.Soft(primary_hue="orange")
)

demo.launch()