import gradio as gr
import torch
from transformers import pipeline
import base64
import os

def get_logo_html():
    if os.path.exists("logo.png"):
        with open("logo.png", "rb") as f:
            b64 = base64.b64encode(f.read()).decode("utf-8")
        img = f'<img src="data:image/png;base64,{b64}" style="width:64px;height:64px;border-radius:12px;box-shadow:0 0 24px #ff6a0055;">'
    else:
        img = ""
    return f"""
    <div class="braingpt-header">
        {img}
        <div>
            <div class="braingpt-title">BrainGPT</div>
            <div class="braingpt-sub">NEURAL ENGINE ONLINE - POWERED BY QWEN2</div>
        </div>
    </div>
    """

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Rajdhani:wght@400;600;700&family=Share+Tech+Mono&display=swap');

*, *::before, *::after { box-sizing: border-box; }

body, .gradio-container {
    background: #0a0a0a !important;
    font-family: 'Rajdhani', sans-serif !important;
    color: #e8e0d0 !important;
}

.braingpt-header {
    display: flex;
    align-items: center;
    gap: 18px;
    padding: 18px 24px 10px 24px;
    border-bottom: 1px solid #2a1a00;
    background: linear-gradient(180deg, #110900 0%, #0a0a0a 100%);
    margin-bottom: 8px;
}

.braingpt-title {
    font-family: 'Rajdhani', sans-serif;
    font-size: 2rem;
    font-weight: 700;
    letter-spacing: 3px;
    background: linear-gradient(90deg, #ff6a00, #ffb347, #ff6a00);
    background-size: 200%;
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    animation: shimmer 3s linear infinite;
    text-transform: uppercase;
}

.braingpt-sub {
    font-family: 'Share Tech Mono', monospace;
    font-size: 0.72rem;
    color: #664422;
    letter-spacing: 2px;
    margin-top: 2px;
}

@keyframes shimmer {
    0% { background-position: 0% }
    100% { background-position: 200% }
}

.chatbot {
    background: #0d0d0d !important;
    border: 1px solid #1f1208 !important;
    border-radius: 12px !important;
    font-family: 'Rajdhani', sans-serif !important;
    font-size: 1.05rem !important;
}

.message.user {
    background: linear-gradient(135deg, #1a0d00, #2a1500) !important;
    border: 1px solid #3d1f00 !important;
    border-radius: 12px 12px 2px 12px !important;
    color: #ffcc88 !important;
    font-weight: 600 !important;
}

.message.bot {
    background: linear-gradient(135deg, #0f0f0f, #161208) !important;
    border: 1px solid #2b1f00 !important;
    border-radius: 12px 12px 12px 2px !important;
    color: #e8d8b8 !important;
}

.thinking-bar {
    display: none;
    align-items: center;
    gap: 10px;
    padding: 8px 16px;
    margin: 4px 0 8px 0;
    background: linear-gradient(90deg, #1a0d00, #0d0d0d);
    border-left: 3px solid #ff6a00;
    border-radius: 0 8px 8px 0;
    font-family: 'Share Tech Mono', monospace;
    font-size: 0.8rem;
    color: #ff8c33;
    letter-spacing: 1px;
}

.thinking-bar.active { display: flex; }

.thinking-dots span {
    display: inline-block;
    width: 6px;
    height: 6px;
    background: #ff6a00;
    border-radius: 50%;
    margin: 0 2px;
    animation: bounce-dot 1.2s infinite;
}
.thinking-dots span:nth-child(2) { animation-delay: 0.2s; }
.thinking-dots span:nth-child(3) { animation-delay: 0.4s; }

@keyframes bounce-dot {
    0%, 80%, 100% { transform: translateY(0); opacity: 0.4; }
    40% { transform: translateY(-6px); opacity: 1; }
}

#msg-input textarea {
    background: #111 !important;
    border: 1px solid #2b1800 !important;
    border-radius: 10px !important;
    color: #ffcc88 !important;
    font-family: 'Rajdhani', sans-serif !important;
    font-size: 1rem !important;
    caret-color: #ff6a00 !important;
    resize: none !important;
    padding: 12px 14px !important;
    transition: border-color 0.2s;
}

#msg-input textarea:focus {
    border-color: #ff6a00 !important;
    box-shadow: 0 0 12px #ff6a0033 !important;
    outline: none !important;
}

#send-btn, #clear-btn {
    font-family: 'Rajdhani', sans-serif !important;
    font-weight: 700 !important;
    letter-spacing: 1.5px !important;
    border-radius: 10px !important;
    border: none !important;
    cursor: pointer !important;
    transition: all 0.18s ease !important;
    text-transform: uppercase !important;
    font-size: 0.9rem !important;
    padding: 12px 22px !important;
}

#send-btn {
    background: linear-gradient(135deg, #ff6a00, #cc4400) !important;
    color: #fff !important;
    box-shadow: 0 0 16px #ff6a0044 !important;
    min-width: 110px !important;
}

#send-btn:hover {
    background: linear-gradient(135deg, #ff8c33, #ff5500) !important;
    box-shadow: 0 0 28px #ff6a0077 !important;
    transform: translateY(-1px) !important;
}

#send-btn.thinking {
    background: linear-gradient(135deg, #2a1500, #1a0d00) !important;
    color: #ff8c33 !important;
    pointer-events: none !important;
}

#clear-btn {
    background: #161208 !important;
    color: #664422 !important;
    border: 1px solid #2b1800 !important;
}

#clear-btn:hover {
    background: #1f1208 !important;
    color: #ff8c33 !important;
    border-color: #ff6a00 !important;
}

.examples-table td {
    background: #111 !important;
    border: 1px solid #1f1200 !important;
    color: #cc7733 !important;
    font-family: 'Rajdhani', sans-serif !important;
    border-radius: 6px !important;
    cursor: pointer !important;
    transition: background 0.15s, color 0.15s !important;
}

.examples-table td:hover {
    background: #1f1000 !important;
    color: #ff8c33 !important;
    border-color: #ff6a00 !important;
}

::-webkit-scrollbar { width: 5px; }
::-webkit-scrollbar-track { background: #0a0a0a; }
::-webkit-scrollbar-thumb { background: #2b1800; border-radius: 4px; }
::-webkit-scrollbar-thumb:hover { background: #ff6a00; }
"""

JS = """
() => {
    function setup() {
        const sendBtn = document.getElementById('send-btn');
        const bar = document.getElementById('thinking-bar');
        if (!sendBtn || !bar) return;
        new MutationObserver(() => {
            const loading = document.querySelector('.generating') !== null
                         || document.querySelector('.eta-bar') !== null;
            if (loading) {
                sendBtn.classList.add('thinking');
                const s = sendBtn.querySelector('span');
                if (s) s.innerText = 'Thinking...';
                bar.classList.add('active');
            } else {
                sendBtn.classList.remove('thinking');
                const s = sendBtn.querySelector('span');
                if (s) s.innerText = 'Enviar';
                bar.classList.remove('active');
            }
        }).observe(document.body, { childList: true, subtree: true });
    }
    setTimeout(setup, 1500);
}
"""

pipe = pipeline(
    "text-generation",
    model="Qwen/Qwen2-1.5B-Instruct",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto",
)

def chat(message, history):
    if not message or not message.strip():
        return ""
    try:
        if len(message) > 1000:
            message = message[:1000] + "... [texto cortado]"

        system_msg = (
            "Voce e BrainGPT, uma inteligencia artificial poderosa, direta e precisa. "
            "Responda sempre em portugues de forma clara e concisa. "
            "Nunca invente factos. Se nao souber, diga claramente."
        )

        msgs = [{"role": "system", "content": system_msg}]
        for u, b in history[-6:]:
            if u: msgs.append({"role": "user", "content": u})
            if b: msgs.append({"role": "assistant", "content": b})
        msgs.append({"role": "user", "content": message})

        prompt = pipe.tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)

        if len(prompt) > 3500:
            return "Conversa muito longa. Clica em Limpar e recomeça."

        result = pipe(
            prompt,
            max_new_tokens=280,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            pad_token_id=pipe.tokenizer.eos_token_id,
            truncation=True,
        )

        resposta = result[0]["generated_text"].split("<|im_start|>assistant\n")[-1].strip()
        resposta = resposta.replace("<|im_end|>", "").strip()
        return resposta if resposta else "Resposta vazia. Tenta reformular."

    except torch.cuda.OutOfMemoryError:
        return "Memoria GPU esgotada. Clica em Limpar e tenta novamente."
    except Exception as e:
        return f"Erro: {str(e)[:150]}"


def user_submit(message, history):
    if not message.strip():
        return "", history
    return "", history + [[message, None]]

def bot_respond(history):
    if not history or history[-1][1] is not None:
        return history
    history[-1][1] = chat(history[-1][0], history[:-1])
    return history


with gr.Blocks(title="BrainGPT") as demo:

    gr.HTML(get_logo_html())

    gr.HTML("""
    <div id="thinking-bar" class="thinking-bar">
        <div class="thinking-dots"><span></span><span></span><span></span></div>
        <span>THINKING...</span>
    </div>
    """)

    chatbot = gr.Chatbot(label="", height=480, bubble_full_width=False)

    with gr.Row():
        msg = gr.Textbox(
            placeholder="Escreve a tua mensagem...",
            show_label=False,
            lines=1,
            max_lines=5,
            elem_id="msg-input",
            scale=8,
        )
        send_btn = gr.Button("Enviar", elem_id="send-btn", variant="primary", scale=2)
        clear_btn = gr.Button("Limpar", elem_id="clear-btn", variant="secondary", scale=1)

    gr.Examples(
        examples=["Quem es tu?", "Explica IA em 3 frases", "Como funciona o Bitcoin?", "5 dicas de produtividade"],
        inputs=msg,
        label="Exemplos",
    )

    msg.submit(user_submit, [msg, chatbot], [msg, chatbot], queue=False).then(bot_respond, chatbot, chatbot)
    send_btn.click(user_submit, [msg, chatbot], [msg, chatbot], queue=False).then(bot_respond, chatbot, chatbot)
    clear_btn.click(lambda: ([], ""), outputs=[chatbot, msg], queue=False)

demo.queue(max_size=5).launch(css=CSS, js=JS)