import os
import pinecone
from langchain.vectorstores import Pinecone
from langchain.embeddings import OpenAIEmbeddings
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from dotenv import load_dotenv, find_dotenv
load_dotenv(find_dotenv(), override=True)

import os
import pinecone
import gradio as gr
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.vectorstores import Pinecone
from langchain.embeddings.openai import OpenAIEmbeddings

from dotenv import load_dotenv, find_dotenv
load_dotenv(find_dotenv(), override=True)

PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
PINECONE_ENV = os.environ.get('PINECONE_ENV')

embeddings = OpenAIEmbeddings()

pinecone.init(
    api_key=PINECONE_API_KEY,
    environment=PINECONE_ENV
)

index_name = "okean"
index = pinecone.Index(index_name)

llm = ChatOpenAI(model='gpt-4-1106-preview', temperature=0)

template="""Assistente é uma IA que tira dúvidas de um manual e documentação. Os documentos tem conteudo em ingles e portugues mas responda sempre em português.
    Assistente elabora repostas simplificadas, com base no contexto fornecido.
    Assistente fornece referências extraídas do contexto abaixo. 
    Caso não consiga encontrar no contexto abaixo ou caso a pergunta não esteja relacionada do contexto do manual, 
    diga apenas 'Eu não sei!'
    Pergunta: {query}

    Contexto: {context}
    Referência: [Página: {page}, Fonte: {source}]
"""

prompt = PromptTemplate(
    template=template,
    input_variables=["query", "context", "page", "source"]
)

historico = []

def clear_history():
    global historico
    historico = []
    return "Histórico limpo!", "Histórico de perguntas e respostas aqui..."


def search(query):
    global historico
    # Adiciona a pergunta ao histórico
    historico.append(f"Pergunta: {query}")

    # Realiza a busca por similaridade
    docsearch = Pinecone.from_existing_index(index_name, embeddings)
    docs = docsearch.similarity_search(query, k=3)

    first_doc = docs[0]
    page_ref = first_doc.metadata['page']
    source_ref = first_doc.metadata['source']

    # Concatena o conteúdo das páginas com o histórico e inclui os metadados
    context = ' '.join(historico)
    for doc in docs:
        print(doc)
        context += f" {doc.page_content} [Referência: Página - {doc.metadata['page']}, Fonte - {doc.metadata['source']}]"


    # Chama a cadeia LLM com o novo contexto
    resp = LLMChain(prompt=prompt, llm=llm)    
    answer = resp.run(query=query, context=context, page=page_ref, source=source_ref)

    # Formatar a resposta para incluir as informações de página e fonte
    formatted_answer = f"{answer}\n\nReferência: Página - {page_ref}, Fonte - {source_ref}"


    # Adiciona a resposta ao histórico
    historico.append(f"Resposta: {formatted_answer}")

    # Retorne tanto a resposta quanto o histórico atualizado
    return formatted_answer, "\n".join(historico)



with gr.Blocks(title="Chatbot Inteligente", theme=gr.themes.Soft()) as ui:
    gr.Markdown("# Sou uma IA que tem o manual OKEAN como base de conhecimento")
    query = gr.Textbox(label='Faça a sua pergunta:', placeholder="EX: como funcionam o sistema de propulsão?")
    text_output = gr.Textbox(label="Resposta")
    historico_output = gr.Textbox(label="Histórico", value="Histórico de perguntas e respostas aqui...")
    btn = gr.Button("Perguntar")
    clear_btn = gr.Button("Limpar Histórico")  # Novo botão para limpar o histórico
    btn.click(fn=search, inputs=query, outputs=[text_output, historico_output])
    clear_btn.click(fn=clear_history, inputs=[], outputs=[text_output, historico_output])  # Vincula a nova função ao botão
ui.launch(debug=True, share=True)