import json

import faiss
import numpy as np
import streamlit as st
from groq import Groq
from sentence_transformers import SentenceTransformer

from streamlit_extras.stylable_container import stylable_container

st.set_page_config(page_title='Psicopatologia Geral - Jaspers', layout='wide', page_icon='🧠')

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

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

with stylable_container(
    key="banner",
    css_styles="""
    img {
        width: 1700px;
        height: 280px;
        overflow: hidden;
        position: relative;
        object-fit: cover;
        border-radius: 20px; /* Adiciona bordas arredondadas */
        mask-image: linear-gradient(to bottom, rgba(0, 0, 0, 1), rgba(0, 0, 0, 0));
        -webkit-mask-image: linear-gradient(to bottom, rgba(0, 0, 0, 1), rgba(0, 0, 0, 0)); /* For Safari */
    }
    """,
):
    st.image("25.jpg")


# Inicialização do modelo de embedding
@st.cache_resource
def load_embedding_model():
    return SentenceTransformer('all-MiniLM-L6-v2')

model = load_embedding_model()

client = Groq(
    api_key=st.secrets["GROQ_API_KEY"],
)

@st.cache_resource
def load_index_and_embeddings(index_file: str, embeddings_file: str):
    index = faiss.read_index(index_file)
    embeddings = np.load(embeddings_file)
    return index, embeddings


def search(query: str, index, embeddings: np.ndarray, chunks: list, k: int = 5):
    query_vector = model.encode([query])
    distances, indices = index.search(query_vector.astype('float32'), k)
    return [chunks[i] for i in indices[0]]


def query_groq(prompt, client):
    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": prompt,
            }
        ],
        model="mixtral-8x7b-32768",
        temperature=0.5,
        max_tokens=2500,
    )
    return chat_completion.choices[0].message.content

# Função para converter chunk em texto
def chunk_to_text(chunk):
    if isinstance(chunk, str):
        return chunk
    elif isinstance(chunk, list):
        return " ".join(map(str, chunk))
    else:
        return str(chunk)


# Carregar índice, embeddings e chunks
@st.cache_resource
def load_data():
    index, embeddings = load_index_and_embeddings("index.faiss", "embeddings.npy")
    with open("chunks.json", "r") as f:
        chunks = json.load(f)
    return index, embeddings, chunks

index, embeddings, chunks = load_data()

# Interface Streamlit
st.title("Dr. Pers - Psicopatologia Geral - K. Jaspers")

# Inicializar histórico do chat
if "groq_chat_history" not in st.session_state:
    st.session_state.groq_chat_history = []

# Exibição do histórico do chat
for message in st.session_state.groq_chat_history:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Input do usuário
user_message = st.chat_input("Digite sua pergunta sobre Psicopatologia Geral:")

if user_message:
    # Adiciona a mensagem do usuário ao histórico
    st.session_state.groq_chat_history.append({"role": "user", "content": user_message})
    
    # Exibe a mensagem do usuário
    with st.chat_message("user"):
        st.markdown(user_message)
    
    # Realiza a busca e gera a resposta
    relevant_chunks = search(user_message, index, embeddings, chunks)
    context = "\n".join(chunk_to_text(chunk) for chunk in relevant_chunks)
    
    prompt = f"""Você é um velho psiquiatra, Dr. Pers, que ajudou a transcrever o livro Psicopatologia geral de Karl Jaspers.
    Sua função é auxiliar jovens psiquiatras e interessados em psicopatologia a entender melhor esta grande obra.
    Com base no seguinte contexto, que tem esta obra completa, responda à pergunta do usuário. Seja sempre direto 
    e elabore bem as respostas, porém caso o usuário queira falar sobre qualquer outro assunto, responda:
    "Não tenho conhecimento concreto para tratar de outros temas senão da Psicopatologia geral." Seja sempre cortês,
    e sempre dê suas respostas em português do Brasil.

    Contexto:
    {context}

    Pergunta: {user_message}

    Resposta:"""
    
    # Gera a resposta
    response = query_groq(prompt, client)
    
    # Adiciona a resposta ao histórico
    st.session_state.groq_chat_history.append({"role": "assistant", "content": response})
    
    # Exibe a resposta
    with st.chat_message("assistant"):
        st.markdown(response)

# Informações adicionais
st.sidebar.title("Informações")

with stylable_container(
    key="jas",
    css_styles="""
    img {
        position: relative;
        object-fit: cover;
        border-radius: 18px;
    }
    """,
):
    st.sidebar.image("jas1.jpg")

def clear_chat_history():
    """
    Clears the chat history and resets the initial analysis in the session state.
    This function clears the chat history and resets the initial analysis in the session state.
    """
    st.session_state.groq_chat_history = []

if st.sidebar.button("Reiniciar a conversa", type="primary"):
    clear_chat_history()

st.sidebar.image('ap.jpg', width=150)
st.sidebar.markdown(
    """
    #### Informações:
    ###### - Psicopatologia Geral, I e II - Jaspers 
    ##### - Acesso livre
    ### Links:
    ##### - [Obsidian - Dr Guilherme](http://dr-guilhermeapolinario.com) 🌎
    ##### - [GitHub - Dr Guilherme](http://dr-guilhermeapolinario.com) 🌎
    """
)