import json
import os

import faiss
import numpy as np
import streamlit as st
from openai import OpenAI
from sentence_transformers import SentenceTransformer

from prompts_template import prompt

# Initialize the messages list in the session state
if "messages" not in st.session_state:
    st.session_state.messages = []

STYLE = "static/style.css"
with open(STYLE, "r", encoding="utf-8") as f:
    st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)

# STYLE = "static/styles.html"
# with open(STYLE, "r", encoding="utf-8") as f:
#     st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)

client = OpenAI(api_key=os.environ['OPENAI_API'])
model = SentenceTransformer("all-MiniLM-L6-v2")

# índice FAISS
json_file = "scr/faiss/chunks .json"
embeddings_file = "scr/faiss/embeddings.npy"
index_file = "scr/faiss/faiss_index.index"
text_chunks_file = "scr/faiss/text_chunks.npy"

with open(json_file, "r") as file:
    chunks = json.load(file)

embeddings = np.load(embeddings_file)
index = faiss.read_index(index_file)
text_chunks = np.load(text_chunks_file)

with st.sidebar:
    openai_api_key = st.text_input(
        "OpenAI API Key", key="chatbot_api_key", type="password"
    )
    st.markdown(
        "[Pegue aqui sua chave OpenAI API](https://platform.openai.com/account/api-keys)"
    )
if st.sidebar.button("Limpar Conversa"):
    st.session_state.messages = []
    st.rerun()

def retrieve(query, k=3):
    """
    Retrieves the most similar document chunks to a given query.
    """
    query_embedding = model.encode([query])
    distances, indices = index.search(query_embedding, k)
    return [chunks[idx] for idx in indices[0]]


def generate_response(query, context):
    """
    Generates a response to a given query based on a provided context.
    """
    prompt_query = (
        f"Com base nos seguintes dados: {context}\n\nPergunta: {query}\n\nResposta:"
    )
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role":
                "system",
                "content":
                "Você é um assistente especializado em análise de encaminhamentos médicos.",
            },
            {
                "role": "user",
                "content": prompt_query
            },
        ],
    )
    return response.choices[0].message.content

co1, co2, co3 = st.columns([1.5, 0.4, 3])
with co1:
    st.write('')
with co2:
    st.image('icons/icon.svg', width=80)

with co3:
    st.title("Encaminhamento Médico")

col1, col2, col3, col4, col5 = st.columns([1, 3, 0.5, 3, 1])

fem= 'icons/iconF.svg'
with col1:
    st.write('')

with col2:
    st.header(" :black_nib: Entrada de Dados")
    st.markdown("#### ")
    # Campos para preenchimento do relatório
    cl1, cl2, cl3 = st.columns([2, 1, 2])
    with cl1:
        idade = st.text_input("###### IDADE", value="")
        sexo_opcao = st.radio("Sexo", ("Masculino", "Feminino"),
              index=0)
        sexo = (sexo_opcao == "Masculino")
    with cl2:
        st.write("")
    with cl3:
        st.write("Condições")
        has = st.checkbox("HAS", value=False)
        dm = st.checkbox("DM TIPO 2", value=False)
        tabaco = st.checkbox("TABAGISTA?", value=False)
        alcool = st.checkbox("ALCOOLISTA?", value=False)

    comorbidades = st.text_area("Outras Comorbidades e medicamentos em uso")
    motivo = st.text_area("Motivo do Encaminhamento")
    historia_clinica = st.text_area("História Clínica Resumida")
    exame_fisico = st.text_area(
        "Exame Físico Relevante",
        "EF:\nBEG, hidratado, s/ edema\nAR: MVUA s/ RA | AC: RCR 2t bnf",
    )
    exames_complementares = st.text_area("Exames Complementares")
    hipotese_diagnostica = st.text_input("Hipótese Diagnóstica")
    justificativa = st.text_area("Justificativa do Encaminhamento")

    # Prioridade
    prioridade = st.selectbox("Prioridade",
                              ["Selecione...", "P0", "P1", "P2", "P3"])

    especialidades = [
        "Cardiologia",
        "Dermatologia",
        "Endocrinologia",
        "Gastroenterologia",
        "Geriatria",
        "Hematologia",
        "Infectologia",
        "Medicina de Família e Comunidade",
        "Nefrologia",
        "Neurologia",
        "Obstetrícia",
        "Oftalmologia",
        "Oncologia",
        "Ortopedia",
        "Otorrinolaringologia",
        "Pediatria",
        "Pneumologia",
        "Psiquiatria",
        "Reumatologia",
        "Urologia",
    ]

    especialidade_selecionada = st.multiselect(
        "Selecione a(s) Especialidade(s):", especialidades)

with col3:
    st.write("")

with col4:
    st.header(" 🤖  Relatório Gerado por IA ")

    # Botão para gerar relatório
    if st.button("Gerar Relatório"):
        user_prompt = f"""
                    Relatório de Encaminhamento Médico

                    1. Identificação do Paciente
                    Idade: {idade}, {sexo}

                    2. Fatores de Risco e Comorbidades
                    HAS: {'Sim' if has else 'Não'}
                    DM TIPO 2: {'Sim' if dm else 'Não'}
                    TABAGISTA: {'Sim' if tabaco else 'Não'}
                    ALCOOLISTA: {'Sim' if alcool else 'Não'}
                    Outras Comorbidades e medicamentos em uso: {comorbidades}

                    3. Motivo do Encaminhamento
                    {motivo}

                    4. História Clínica Resumida
                    {historia_clinica}

                    5. Exame Físico Relevante
                    {exame_fisico}

                    6. Exames Complementares
                    {exames_complementares}

                    7. Hipótese Diagnóstica
                    {hipotese_diagnostica}

                    8. Justificativa do Encaminhamento
                    {justificativa}

                    9. Prioridade
                    {prioridade}

                    10. Especialidade(s) de Destino
                    {', '.join(especialidade_selecionada)}

                    Com base nessas informações, elabore um relatório de
                    encaminhamento médico conciso, mas informativo.
                    Certifique-se de:

                    1. Manter uma linguagem profissional e clara.
                    2. Destacar os pontos mais relevantes para a especialidade
                    de destino.
                    3. Incluir apenas informações pertinentes ao encaminhamento.
                    4. Justificar claramente a necessidade do encaminhamento e a
                    prioridade atribuída.
                    5. Limitar o relatório a no máximo 300 palavras.

                    Por favor, gere o relatório mantendo a estrutura fornecida,
                    mas adaptando o conteúdo para ser mais fluido e coeso.
                    """
        st.text_area("Prompt gerado:", user_prompt, height=600)

        completion = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{
                "role": "system",
                "content": prompt() + user_prompt
            }],
            temperature=0.4,
            max_tokens=1500,
        )
        resposta = completion.choices[0].message.content
        st.text_area("Resposta da IA:", resposta, height=800)
        st.write(
            "O relatório gerado pela IA será exibido aqui após o processamento dos dados inseridos."
        )

#---------------------------------------------------------------

# Interface do chat
st.title("Chat Médico Baseado em Query")

# Exibe as mensagens já enviadas
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Campo de entrada de texto para o usuário
if prompt := st.chat_input("Digite sua pergunta"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Processa a query e gera a resposta
    retrieved_context = retrieve(prompt)
    response = generate_response(prompt, retrieved_context)

    # Adiciona a resposta do bot
    st.session_state.messages.append({
        "role": "assistant",
        "content": response
    })
    with st.chat_message("assistant"):
        st.markdown(response)
with col5:
    st.write('')