import os
from flask import Flask, render_template, request, jsonify, redirect, url_for, flash, session
from werkzeug.utils import secure_filename
from dotenv import load_dotenv

# Importações do LangChain
from langchain import hub
from langchain_chroma import Chroma
from langchain_community.document_loaders import PyPDFLoader
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

# Carrega variáveis de ambiente (sua chave da OpenAI)
load_dotenv()

# --- Configuração da Aplicação Flask ---
app = Flask(__name__)
app.secret_key = os.urandom(24) # Necessário para usar 'session' e 'flash'

# --- Configuração de Pastas ---
UPLOAD_FOLDER = 'uploads'
VECTOR_STORE_PATH = 'vector_store'
ALLOWED_EXTENSIONS = {'pdf', 'txt', 'doc', 'docx'}
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(VECTOR_STORE_PATH, exist_ok=True)

# --- Configuração do LangChain (Global) ---
# Inicializamos os componentes que não dependem dos documentos
os.environ["OPENAI_API_KEY"] = os.getenv('OPENAI_API_KEY') # Use a sua chave aqui ou em um arquivo .env

model = ChatOpenAI(model='gpt-4')
embedding = OpenAIEmbeddings()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
prompt = hub.pull('rlm/rag-prompt')

# Variável global para o nosso RAG chain
rag_chain = None

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def initialize_rag_chain():
    """
    Inicializa ou atualiza o Vector Store e o RAG Chain.
    Esta função irá varrer a pasta 'uploads' e processar qualquer arquivo novo.
    """
    global rag_chain

    # Carrega documentos da pasta 'uploads'
    loaded_documents = []
    for filename in os.listdir(app.config['UPLOAD_FOLDER']):
        if allowed_file(filename):
            file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            # Para simplificar, vamos usar PyPDFLoader. Para outros tipos, você precisaria de outros loaders.
            if filename.endswith('.pdf'):
                loader = PyPDFLoader(file_path)
                loaded_documents.extend(loader.load())
    
    if not loaded_documents:
        print("Nenhum documento encontrado para processar.")
        return

    # Divide os documentos em chunks
    chunks = text_splitter.split_documents(documents=loaded_documents)
    
    # Cria ou carrega o Vector Store persistente
    vector_store = Chroma.from_documents(
        documents=chunks,
        embedding=embedding,
        persist_directory=VECTOR_STORE_PATH
    )

    retriever = vector_store.as_retriever()

    # Cria a RAG chain
    rag_chain = (
        {
            'context': retriever,
            'question': RunnablePassthrough(),
        }
        | prompt
        | model
        | StrOutputParser()
    )
    print("RAG Chain inicializado com sucesso!")

# --- Rotas da Aplicação ---

@app.route('/')
def index():
    """Página principal do chat."""
    # Inicializa o histórico do chat na sessão se não existir
    if 'chat_history' not in session:
        session['chat_history'] = [
            {'sender': 'assessor', 'message': 'Olá! Como posso te ajudar hoje?'}
        ]
    return render_template('index.html', chat_history=session['chat_history'])

@app.route('/ask', methods=['POST'])
def ask():
    """Endpoint para receber perguntas do usuário via AJAX."""
    if not rag_chain:
        return jsonify({'answer': 'Desculpe, a base de conhecimento ainda não foi inicializada. Por favor, adicione um documento.'})

    user_question = request.json.get('question')
    if not user_question:
        return jsonify({'error': 'Nenhuma pergunta fornecida.'}), 400

    # Adiciona a pergunta do usuário ao histórico
    session['chat_history'].append({'sender': 'user', 'message': user_question})

    # Obtém a resposta do RAG chain
    response = rag_chain.invoke(user_question)

    # Adiciona a resposta do assessor ao histórico
    session['chat_history'].append({'sender': 'assessor', 'message': response})
    session.modified = True # Marca a sessão como modificada

    return jsonify({'answer': response})

@app.route('/knowledge_base')
def knowledge_base():
    """Página para gerenciar a base de conhecimento."""
    files = os.listdir(app.config['UPLOAD_FOLDER'])
    return render_template('knowledge_base.html', files=files)

@app.route('/upload', methods=['POST'])
def upload_file():
    """Endpoint para upload de novos documentos."""
    if 'file' not in request.files:
        flash('Nenhum arquivo selecionado')
        return redirect(url_for('knowledge_base'))
    
    file = request.files['file']
    if file.filename == '':
        flash('Nenhum arquivo selecionado')
        return redirect(url_for('knowledge_base'))

    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(file_path)
        flash(f'Arquivo "{filename}" enviado com sucesso! Re-processando a base de conhecimento...')
        
        # Re-inicializa o RAG chain para incluir o novo documento
        initialize_rag_chain()
        
        return redirect(url_for('knowledge_base'))
    else:
        flash('Tipo de arquivo não permitido.')
        return redirect(url_for('knowledge_base'))

# --- Inicialização ---
if __name__ == '__main__':
    # Na primeira vez que o servidor iniciar, cria o RAG chain
    initialize_rag_chain()
    # Para produção, use um servidor WSGI como Gunicorn ou Waitress
    app.run(debug=True)