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

# ==========================================
# 1. MODEL CONFIGURATION
# ==========================================

MODEL_ID = "DavidBazaldua/llama3_finetuned_transformes"
DEVICE = "cpu"
DTYPE = torch.float32

print(f"Loading model: {MODEL_ID}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=DTYPE,
    low_cpu_mem_usage=True,
)
model.to(DEVICE)
model.eval()
print("Model loaded successfully.")

DEFAULT_SYSTEM_PROMPT = (
    "You are a helpful, precise AI assistant. "
    "Always answer as briefly as possible. "
    "Respond in English unless the user asks otherwise."
)

# ==========================================
# 2. CHAT LOGIC
# ==========================================

def build_prompt(system_prompt, context, history, user_message):
    messages = []
    if system_prompt and system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt})
    if context and context.strip():
        messages.append({
            "role": "system",
            "content": f"Use the following context if relevant:\n{context}"
        })
    for user, assistant in history:
        messages.append({"role": "user", "content": user})
        messages.append({"role": "assistant", "content": assistant})
    messages.append({"role": "user", "content": user_message})

    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True,
    )
    return prompt

def chat_logic(message, history, system_prompt, context, max_tokens, temperature, top_p):
    if not message: return "", history
    if history is None: history = []
    if not system_prompt: system_prompt = DEFAULT_SYSTEM_PROMPT

    max_tokens = int(min(max_tokens, 1024))
    prompt = build_prompt(system_prompt, context, history, message)
    inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(DEVICE)

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

    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
    if decoded.startswith(prompt):
        answer = decoded[len(prompt):].strip()
    else:
        answer = decoded.split("assistant\n")[-1].strip()

    history.append((message, answer))
    return "", history

# ==========================================
# 3. UI DESIGN (FINAL FIX)
# ==========================================

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;600;700&display=swap');

/* --- GLOBAL RESET --- */
:root, .dark, body, .gradio-container {
    --body-text-color: #000000 !important;
    --body-background-fill: transparent !important;
    --background-fill-primary: transparent !important;
    --background-fill-secondary: transparent !important;
    --block-background-fill: transparent !important;
    --block-border-color: transparent !important;
    --input-background-fill: #ffffff !important;
    
    /* COLORS */
    --primary-500: #000000 !important;
    --primary-600: #000000 !important;
    --secondary-500: #000000 !important;
}

body {
    font-family: 'Inter', sans-serif !important;
    background: radial-gradient(circle at 10% 10%, #ffeef8 0%, #edf6ff 50%, #f7f0ff 90%) !important;
    color: #000000 !important;
    margin: 0; padding: 0;
    min-height: 100vh;
}

/* HIDE GRADIO FOOTER (CSS is enough) */
footer {
    display: none !important;
    visibility: hidden !important;
}

h1 { font-weight: 700 !important; color: #000000 !important; letter-spacing: -0.02em; }
p, span, label { color: #374151 !important; }

/* --- LEFT PANEL --- */
.config-card {
    background: rgba(255, 255, 255, 0.4) !important; 
    border: 1px solid rgba(255, 255, 255, 0.5) !important;
    backdrop-filter: blur(10px);
    border-radius: 24px;
    padding: 24px !important;
}

.config-card div, .config-card .gr-form, .config-card .gr-box, 
.config-card .gr-group, .config-card .gr-block, .config-card .gr-panel {
    background: transparent !important;
    border: none !important;
    box-shadow: none !important;
}

.config-card label span, 
.config-card .label-wrap span,
.config-card .block-title {
    background-color: transparent !important;
    background: transparent !important;
    color: #000000 !important;
    border: none !important;
    font-weight: 700 !important;
    font-size: 0.85rem !important;
    margin-bottom: 4px !important;
    padding: 0 !important;
}

.config-card textarea, 
.config-card input[type="text"], 
.config-card input[type="number"] {
    background-color: #ffffff !important;
    color: #000000 !important;
    border: 1px solid #e5e7eb !important;
    border-radius: 12px !important;
    box-shadow: 0 1px 2px rgba(0,0,0,0.05) !important;
}

input[type=range] {
    accent-color: #000000 !important;
    filter: grayscale(100%) brightness(0) !important;
}

/* --- RIGHT PANEL: CHAT --- */
.chat-window {
    background: rgba(255, 255, 255, 0.95) !important;
    border: 1px solid rgba(255, 255, 255, 0.9);
    backdrop-filter: blur(20px);
    border-radius: 24px;
    box-shadow: 0 20px 50px rgba(0,0,0,0.05);
    overflow: hidden;
    padding: 0 !important;
}

.window-header {
    display: flex; gap: 8px; align-items: center;
    padding: 18px 24px 10px 24px;
}
.dot { width: 12px; height: 12px; border-radius: 50%; }
.red { background-color: #ff5f56; }
.yellow { background-color: #ffbd2e; }
.green { background-color: #27c93f; }

#chatbot {
    background: transparent !important;
    border: none !important;
    height: 550px !important;
}

/* --- CHAT BUBBLES --- */
.message.user {
    background-color: #000000 !important;
    color: #ffffff !important;
    border-radius: 18px 18px 4px 18px !important;
    font-weight: 500;
}
.message.user p, .message.user span, .message.user div {
    color: #ffffff !important;
}

.message.bot {
    background-color: #ffffff !important;
    color: #000000 !important;
    border: 1px solid #e5e7eb !important;
    border-radius: 18px 18px 18px 4px !important;
    box-shadow: 0 2px 5px rgba(0,0,0,0.02);
}

/* --- INPUT BOX --- */
.chat-input-area {
    padding: 16px 24px 24px 24px;
    display: flex; gap: 10px; align-items: flex-end;
}

#msg-box textarea {
    background: #ffffff !important;
    padding: 12px 18px !important;
    border: 1px solid #e5e7eb !important;
    color: #000000 !important;
    max-height: 120px !important;
    overflow-y: auto !important;
    border-radius: 24px !important;
}

#send-btn {
    background: transparent !important;
    color: #000000 !important;
    border: none !important;
    box-shadow: none !important;
    font-size: 1.6rem;
    padding: 0 10px !important;
    width: auto !important;
    margin-bottom: 8px;
}
#send-btn:hover {
    transform: translateX(3px);
    color: #333;
}
"""

WINDOW_HEADER = """
<div class="window-header">
    <div class="dot red"></div>
    <div class="dot yellow"></div>
    <div class="dot green"></div>
    <span style="margin-left: 12px; font-size: 0.85rem; color: #6b7280; font-weight: 500;">Iris Assistant</span>
</div>
"""

with gr.Blocks(theme=gr.themes.Soft(), css=CSS, title="Custom AI") as demo:
    
    with gr.Row(elem_id="main-container"):
        
        # --- LEFT PANEL ---
        with gr.Column(scale=4):
            gr.Markdown(
                """
                # MX Iris AI Chatbot
                """
            )
            
            with gr.Group(elem_classes="config-card"):
                gr.Markdown("### ⚙️ Settings")
                
                context_input = gr.Textbox(
                    label="Context (Data)",
                    placeholder="Paste context, notes, or data here...",
                    lines=5
                )
                
                system_prompt_input = gr.Textbox(
                    label="System Prompt",
                    value=DEFAULT_SYSTEM_PROMPT,
                    lines=2
                )
                
                with gr.Row():
                    temp_slider = gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature")
                    tokens_slider = gr.Slider(64, 1024, value=256, step=64, label="Max Tokens")
                
                top_p_slider = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-P")

        # --- RIGHT PANEL ---
        with gr.Column(scale=6):
            with gr.Column(elem_classes="chat-window"):
                gr.HTML(WINDOW_HEADER)
                
                chatbot = gr.Chatbot(
                    elem_id="chatbot",
                    show_label=False,
                    avatar_images=(None, "https://cdn-icons-png.flaticon.com/512/4712/4712027.png"),
                    height=550,
                    layout="bubble"
                )
                
                with gr.Row(elem_classes="chat-input-area"):
                    msg = gr.Textbox(
                        show_label=False,
                        placeholder="Ask a question...",
                        scale=9,
                        elem_id="msg-box",
                        container=False,
                        autofocus=True,
                        lines=1
                    )
                    send_btn = gr.Button("➝", scale=1, elem_id="send-btn")

    inputs_list = [msg, chatbot, system_prompt_input, context_input, tokens_slider, temp_slider, top_p_slider]
    msg.submit(chat_logic, inputs_list, [msg, chatbot])
    send_btn.click(chat_logic, inputs_list, [msg, chatbot])

if __name__ == "__main__":
    # Removed incorrect show_footer argument. CSS handles footer hiding.
    demo.launch(show_api=False)