#!/usr/bin/env python3
import sys
import base64
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

from dotenv import load_dotenv
load_dotenv(Path(__file__).parent / ".env", override=True)

import gradio as gr
from src.labdaps.ingestion.embedder import Embedder
from src.labdaps.chat.session import ChatSession
from src.labdaps.retrieval.vector_store import collection_count

ASSETS_DIR = Path(__file__).parent / "assets"

with open(ASSETS_DIR / "logo.png", "rb") as f:
    LOGO_B64 = base64.b64encode(f.read()).decode()
LOGO_SRC = f"data:image/png;base64,{LOGO_B64}"

with open(ASSETS_DIR / "icon.png", "rb") as f:
    ICON_B64 = base64.b64encode(f.read()).decode()
ICON_SRC = f"data:image/png;base64,{ICON_B64}"

print("[INFO] Inicializando modelo de embeddings...")
embedder = Embedder()
session = ChatSession(embedder)

_last_sources_md = "As fontes aparecerão aqui após cada resposta."


def format_sources(chunks) -> str:
    if not chunks:
        return "_Nenhuma fonte consultada._"
    seen = set()
    lines = []
    for c in chunks:
        key = (c.source_file, c.page_number)
        if key not in seen:
            seen.add(key)
            lines.append(f"- **Sinal de Alerta** -- p. {c.page_number}")
    return "**Fontes consultadas:**\n" + "\n".join(lines)


def respond(message: str, history: list):
    global _last_sources_md
    if not message.strip():
        yield history, _last_sources_md, gr.update()
        return

    # Gradio 5: history e lista de dicts {"role": "user"|"assistant", "content": str}
    new_history = history + [
        {"role": "user", "content": message},
        {"role": "assistant", "content": ""},
    ]
    yield new_history, _last_sources_md, ""

    session.history = [
        {"role": msg["role"], "content": msg["content"]}
        for msg in history
    ]

    partial = ""
    final_chunks = []
    for text_delta, chunks in session.ask(message):
        partial += text_delta
        final_chunks = chunks
        new_history[-1]["content"] = partial
        yield new_history, format_sources(final_chunks), gr.update()

    _last_sources_md = format_sources(final_chunks)
    yield new_history, _last_sources_md, gr.update()


def reset():
    global _last_sources_md
    session.reset()
    _last_sources_md = "As fontes aparecerão aqui após cada resposta."
    return [], _last_sources_md, ""


# Paleta LABDAPS: azul marinho #1C2D6E, verde-limão #8DC63F, fundo escuro #080E1E

LABDAPS_CSS = f"""
/* Loading splash */
#labdaps-splash {{
    position: fixed;
    inset: 0;
    background: #1C2D6E;
    display: flex;
    flex-direction: column;
    align-items: center;
    justify-content: center;
    z-index: 99999;
    transition: opacity 0.5s ease;
}}
#labdaps-splash img {{
    width: 260px;
    margin-bottom: 28px;
    filter: brightness(0) invert(1);
}}
#labdaps-splash .spinner {{
    width: 40px;
    height: 40px;
    border: 4px solid rgba(255,255,255,0.25);
    border-top-color: #8DC63F;
    border-radius: 50%;
    animation: spin 0.8s linear infinite;
}}
@keyframes spin {{ to {{ transform: rotate(360deg); }} }}

/* Global */
body, .gradio-container {{
    background: #080E1E !important;
    color: #E8EBF5 !important;
    font-family: 'Inter', 'Segoe UI', sans-serif !important;
}}

/* Header */
#labdaps-header {{
    background: linear-gradient(135deg, #1C2D6E 0%, #0F1A45 100%);
    padding: 8px 16px;
    border-radius: 10px;
    display: flex;
    align-items: center;
    gap: 16px;
    margin-bottom: 4px;
    border-bottom: 3px solid #8DC63F;
    flex-shrink: 0;
}}
#labdaps-header img {{
    height: 34px;
    filter: brightness(0) invert(1);
}}
#labdaps-header .tagline {{
    color: rgba(255,255,255,0.8);
    font-size: 0.88em;
    margin: 0;
}}
#labdaps-header .book-title {{
    color: #8DC63F;
    font-weight: 700;
    font-size: 0.95em;
    margin: 2px 0 0 0;
}}

/* Status */
#labdaps-status p, #labdaps-status em, #labdaps-status span,
#labdaps-status * {{
    color: #a8b4d0 !important;
    font-size: 13px !important;
    margin: 4px 0 12px 0;
}}

/* Chatbot */
.gradio-chatbot {{
    background: #0D1530 !important;
    border: 1px solid #2D3F7A !important;
    border-radius: 12px !important;
}}
.gradio-chatbot .message.user {{
    background: #ffffff !important;
    color: #111111 !important;
    border-radius: 12px 12px 4px 12px !important;
    border: 1px solid #c8d0e8 !important;
}}
.gradio-chatbot .message.bot {{
    background: #ffffff !important;
    color: #111111 !important;
    border-radius: 12px 12px 12px 4px !important;
    border: 1px solid #c8d0e8 !important;
}}
.gradio-chatbot .message.user *,
.gradio-chatbot .message.bot * {{
    color: #111111 !important;
}}

/* Input */
textarea, input[type=text] {{
    background: #ffffff !important;
    color: #111111 !important;
    border: 1px solid #2D3F7A !important;
    border-radius: 8px !important;
}}
textarea:focus, input[type=text]:focus {{
    border-color: #8DC63F !important;
    box-shadow: 0 0 0 2px rgba(141, 198, 63, 0.25) !important;
}}

/* Buttons */
button.primary {{
    background: #1C2D6E !important;
    border: none !important;
    color: #fff !important;
    border-radius: 8px !important;
    font-weight: 600 !important;
    transition: background 0.2s;
}}
button.primary:hover {{
    background: #8DC63F !important;
    color: #0D1530 !important;
}}
button.secondary {{
    background: transparent !important;
    border: 1px solid #2D3F7A !important;
    color: #8DC63F !important;
    border-radius: 8px !important;
}}
button.secondary:hover {{
    background: rgba(141, 198, 63, 0.1) !important;
}}

/* Sources panel */
#sources-panel {{
    background: #0D1530;
    border: 1px solid #2D3F7A;
    border-radius: 12px;
    padding: 16px;
}}
#sources-panel * {{
    color: #ffffff !important;
}}
#sources-panel p, #sources-panel li,
#sources-panel span, #sources-panel strong,
#sources-panel em {{
    color: #ffffff !important;
    font-size: 13px !important;
}}

/* Labels */
label span {{
    color: #8DC63F !important;
    font-weight: 600 !important;
}}

/* Footer */
#labdaps-footer p {{
    color: rgba(168, 180, 208, 0.7) !important;
    font-size: 12px !important;
    text-align: center;
    margin-top: 8px;
    margin-bottom: 0;
    max-width: 65%;
    margin-left: auto;
    margin-right: auto;
}}
"""

HEAD_HTML = f"""
<link rel="icon" type="image/png" href="{ICON_SRC}">
<div id="labdaps-splash">
  <img src="{LOGO_SRC}" alt="LABDAPS">
  <div class="spinner"></div>
</div>
<script>
  window.addEventListener('load', function() {{
    setTimeout(function() {{
      var splash = document.getElementById('labdaps-splash');
      if (splash) {{
        splash.style.opacity = '0';
        setTimeout(function() {{ splash.style.display = 'none'; }}, 500);
      }}
    }}, 900);

    setTimeout(function() {{
      document.addEventListener('keydown', function(e) {{
        if (e.ctrlKey && e.shiftKey && (e.key === 'L' || e.key === 'l')) {{
          e.preventDefault();
          var btns = document.querySelectorAll('button');
          for (var i = 0; i < btns.length; i++) {{
            if (btns[i].textContent.trim().startsWith('Nova conversa')) {{
              btns[i].click();
              break;
            }}
          }}
        }}
      }}, true);
    }}, 2000);
  }});
</script>
"""

collection_count()
n_chunks = collection_count()
status_msg = (
    f"Base indexada com {n_chunks} trechos do livro. Respostas baseadas exclusivamente no conteúdo do Sinal de Alerta."
    if n_chunks > 0
    else "Execute `python ingest.py` para indexar o livro antes de usar o assistente."
)

with gr.Blocks(
    title="Assistente LABDAPS - Sinal de Alerta",
    css=LABDAPS_CSS,
    head=HEAD_HTML,
    theme=gr.themes.Base(
        primary_hue="blue",
        neutral_hue="slate",
    ),
) as demo:

    gr.HTML(f"""
    <div id="labdaps-header">
      <img src="{LOGO_SRC}" alt="LABDAPS">
      <div>
        <p class="book-title">Sinal de Alerta</p>
        <p class="tagline">Assistente de pesquisa baseado no livro do LABDAPS / FSP-USP</p>
      </div>
    </div>
    """)

    gr.Markdown(status_msg, elem_id="labdaps-status")

    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(label="Chat", type="messages", elem_classes=["gradio-chatbot"], height=420)
            msg_box = gr.Textbox(
                placeholder="Pergunte sobre ML em saúde, predição, modelos... (Enter para enviar)",
                label="Sua pergunta",
                lines=1,
            )
            with gr.Row():
                send_btn = gr.Button("Enviar  (Enter)", variant="primary")
                clear_btn = gr.Button("Nova conversa  (Ctrl+Shift+L)", variant="secondary")
            gr.HTML("""
            <div id="labdaps-footer">
              <p>Este assistente responde <strong style="color:#8DC63F !important;">exclusivamente</strong> com base no livro
              "Sinal de Alerta" (LABDAPS/FSP-USP). Para informações não cobertas, consulte os pesquisadores do laboratório.</p>
            </div>
            """)

        with gr.Column(scale=1):
            gr.Markdown("### <span style='color:#8DC63F'>Fontes consultadas</span>")
            sources_md = gr.Markdown(
                value=_last_sources_md,
                elem_id="sources-panel",
            )

    send_btn.click(fn=respond, inputs=[msg_box, chatbot], outputs=[chatbot, sources_md, msg_box], queue=True)
    msg_box.submit(fn=respond, inputs=[msg_box, chatbot], outputs=[chatbot, sources_md, msg_box], queue=True)
    clear_btn.click(fn=reset, outputs=[chatbot, sources_md, msg_box])


demo.queue()

if __name__ == "__main__":
    demo.launch(server_port=7861, share=False, show_error=True, show_api=False)