import os
import json
import requests
import streamlit as st
from streamlit_extras.stylable_container import stylable_container

# Configuração da API do OpenRouter
OPENROUTER = os.getenv('OPEN_ROUTER_API_KEY')
if not OPENROUTER:
    raise ValueError("API Key não encontrada. Configure a variável de ambiente 'OPEN_ROUTER_API_KEY'.")

url = "https://openrouter.ai/api/v1/chat/completions"
headers = {
    "Authorization": f"Bearer {OPENROUTER}",
    "Content-Type": "application/json"
}

@st.fragment
def chat_imagem():
    with stylable_container(
        key="bot1",
        css_styles="""
        img {
            width: 220px;
            height: 280px;
            overflow: hidden;
            position: relative;
            object-fit: cover;
            border-radius: 24px;
        }
        """
    ):
        st.image("src/Images/bot.png")

cf1, cf2, cf3 = st.columns(3)

with cf1:
    st.title("🤖🧠 - ChatBot UBS Flamengo")
with cf2:
    chat_imagem()
with cf3:
    st.subheader(':blue[🔍📚 Auxílio em temas gerais e administrativos!]')
    st.write("""
    Trata-se do modelo DeepSeek, uma IA com raciocínio, aprimorado. Portanto
    as respostas são mais demoradas, mas mais elaboradas! Após a resposta, 
    deixei em azul, a linha de raciocínio da IA!
    """)

@st.fragment
def clear_all_session_state():
    """
    Remove todas as chaves do st.session_state e reinicia a aplicação.
    """
    for key in list(st.session_state.keys()):
        del st.session_state[key]
    st.rerun()

# Inicialização do histórico do chat
if "ubs_chat_history" not in st.session_state:
    st.session_state.ubs_chat_history = []

for message in st.session_state.ubs_chat_history:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Input do usuário (fora e abaixo do container do histórico)
user_message = st.chat_input("Digite sua pergunta ")

if user_message:
    # Adiciona a mensagem do usuário ao histórico
    st.session_state.ubs_chat_history.append({"role": "user", "content": user_message})
    with st.chat_message("user"):
        st.markdown(user_message)

    try:
        # Contexto e instruções para o assistente
        context = """
Você é um assistente virtual que atua em uma Unidade Básica de Saúde (UBS) chamada UBS Flamengo, em Contagem - Minas Gerais.
Sua função é auxiliar profissionais de saúde e administrativos em diversas tarefas, sempre se comunicando em português brasileiro formal.

1. Comportamento e Comunicação
   - Pergunte inicialmente o nome do usuário e utilize-o de forma respeitosa durante toda a conversa.
   - Mantenha um tom profissional, cortês e empático.
   - Seja claro e compreensível.

2. Principais Responsabilidades
   - Auxiliar na elaboração de relatórios médicos e administrativos.
   - Interpretar e explicar notas técnicas e diretrizes de saúde.
   - Sumarizar textos complexos relacionados à saúde pública.
   - Oferecer insights e análises.
   - Se não souber sobre um tópico específico, admita e sugira a consulta a um profissional qualificado.

3. Formato de Respostas
   - Forneça respostas completas e estruturadas.
   - Para tópicos complexos, apresente uma visão geral concisa seguida de detalhes.
   - Utilize analogias ou exemplos práticos quando relevantes.
        """

        # Constrói o prompt completo combinando o contexto com o histórico da conversa
        full_prompt = context + "\n"
        for msg in st.session_state.ubs_chat_history:
            full_prompt += f"{msg['role']}: {msg['content']}\n"

        payload = {
            "model": "deepseek/deepseek-r1",
            "messages": [
                {"role": "user", "content": full_prompt}
            ],
            "include_reasoning": True,
            "temperature": 0.5,
            "language": "pt-BR"
        }

        response = requests.post(url, headers=headers, data=json.dumps(payload))
        response.raise_for_status()
        json_response = response.json()
        message = json_response['choices'][0]['message']
        content = message.get('content', '')
        reasoning = message.get('reasoning', '')

        # Adiciona a resposta ao histórico
        st.session_state.ubs_chat_history.append({
            "role": "assistant",
            "content": content
        })

        # Exibe a resposta do assistente
        with st.chat_message("assistant"):
            st.markdown(content)
            if reasoning:
                st.info(f"Raciocínio: {reasoning}")

    except Exception as e:
        st.error(f"Ocorreu um erro ao processar a resposta: {str(e)}")
        # Log para análise posterior (opcional)
        with open("error_log.txt", "a") as log_file:
            log_file.write(f"Erro: {str(e)}\n")