import os
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

# ─────────────────────────────────────────────
BASE_MODEL = "Qwen/Qwen3-8B"
LORA_MODEL = "crambrodev/dragonvineAI-qwen3-hytale"

SYSTEM_PROMPT = """You are DragonvineAI — an expert Hytale modding assistant.
You help developers create plugins and mods for Hytale servers.

Key facts about Hytale modding:
- Plugins are written in Java or Kotlin
- Entry point: extend JavaPlugin, implement setup() method
- Manifest file: manifest.json defines plugin metadata
- Build system: Gradle with the hytale-mod plugin
- Commands: extend CommandBase, override executeSync()
- Events: use event listener system to hook into game events
- API package: com.hypixel.hytale.server.core.*

Always provide working, well-commented code examples."""

# ─────────────────────────────────────────────
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)

print("Loading base model...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True,
)

print("Loading LoRA adapter...")
model = PeftModel.from_pretrained(base_model, LORA_MODEL)
model.eval()
print("Model ready!")

# ─────────────────────────────────────────────
def respond(message, history, thinking_mode, max_tokens, temperature):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    messages += history

    prefix = "/think " if thinking_mode else "/no_think "
    messages.append({"role": "user", "content": prefix + message})

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

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=temperature,
            do_sample=temperature > 0,
            pad_token_id=tokenizer.eos_token_id,
        )

    response = tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=True,
    )

    if "<think>" in response and "</think>" in response:
        response = response.split("</think>")[-1].strip()

    return response

# ─────────────────────────────────────────────
with gr.Blocks(title="DragonvineAI — Hytale Modding Assistant", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🐉 DragonvineAI — Hytale Modding Assistant\nAsk anything about creating Hytale plugins and mods!")

    chatbot = gr.Chatbot(height=500, label="Chat", type="messages")

    with gr.Row():
        msg = gr.Textbox(
            placeholder="Ask about Hytale modding... e.g. 'How do I create a custom command?'",
            label="Your question",
            scale=4,
        )
        submit = gr.Button("Send 🚀", scale=1, variant="primary")

    with gr.Accordion("Settings", open=False):
        thinking = gr.Checkbox(label="Thinking mode (slower but smarter)", value=False)
        max_tok  = gr.Slider(128, 1024, value=512, step=64, label="Max tokens")
        temp     = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Temperature")

    gr.Examples(
        examples=[
            "How do I create a simple Hytale plugin with a /hello command?",
            "Show me how to listen to player join events in Hytale",
            "What does a basic manifest.json look like for a Hytale plugin?",
            "How do I register a command in Hytale?",
        ],
        inputs=msg,
    )

    def user_submit(message, history):
        history = history + [{"role": "user", "content": message}]
        return "", history

    def bot_respond(history, thinking, max_tok, temp):
        user_message = history[-1]["content"]
        prev_history = history[:-1]
        response = respond(user_message, prev_history, thinking, max_tok, temp)
        history = history + [{"role": "assistant", "content": response}]
        return history

    submit.click(
        user_submit, inputs=[msg, chatbot], outputs=[msg, chatbot]
    ).then(
        bot_respond, inputs=[chatbot, thinking, max_tok, temp], outputs=chatbot
    )

    msg.submit(
        user_submit, inputs=[msg, chatbot], outputs=[msg, chatbot]
    ).then(
        bot_respond, inputs=[chatbot, thinking, max_tok, temp], outputs=chatbot
    )

demo.launch()