# app.py
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_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI, OpenAIEmbeddings

load_dotenv()

# --- Configuração da Aplicação Flask ---
app = Flask(__name__)
app.secret_key = os.urandom(24)

# --- Configuração de Pastas ---
UPLOAD_FOLDER = 'uploads'
VECTOR_STORE_PATH = 'vector_store'
ALLOWED_EXTENSIONS = {'pdf'} # Simplificado para PDF por enquanto
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 ---
# Inicializamos componentes globais
os.environ["OPENAI_API_KEY"] = os.getenv('OPENAI_API_KEY')
model = ChatOpenAI(model='gpt-4')
embedding = OpenAIEmbeddings()
prompt = hub.pull('rlm/rag-prompt')
rag_chain = None

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

def get_rag_chain():
    """
    Carrega o Vector Store existente e monta o RAG Chain.
    Não processa novos documentos, apenas carrega o que já foi processado.
    """
    global rag_chain
    if rag_chain:
        return rag_chain

    # Verifica se o banco de vetores já existe
    if not os.path.exists(VECTOR_STORE_PATH) or not os.listdir(VECTOR_STORE_PATH):
        print("Vector Store não encontrado ou vazio. Execute o script de processamento primeiro.")
        return None

    # Carrega o Vector Store persistente
    vector_store = Chroma(
        persist_directory=VECTOR_STORE_PATH,
        embedding_function=embedding
    )
    retriever = vector_store.as_retriever()

    rag_chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | model
        | StrOutputParser()
    )
    print("RAG Chain carregado com sucesso!")
    return rag_chain

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

@app.route('/')
def index():
    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():
    chain = get_rag_chain()
    if not chain:
        return jsonify({'answer': 'Desculpe, a base de conhecimento ainda não foi inicializada. Por favor, adicione um documento e aguarde o processamento.'})

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

    session['chat_history'].append({'sender': 'user', 'message': user_question})
    response = chain.invoke(user_question)
    session['chat_history'].append({'sender': 'assessor', 'message': response})
    session.modified = True
    return jsonify({'answer': response})

@app.route('/knowledge_base')
def knowledge_base():
    files = os.listdir(app.config['UPLOAD_FOLDER'])
    return render_template('knowledge_base.html', files=files)

@app.route('/upload', methods=['POST'])
def upload_file():
    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)
        # A MÁGICA ESTÁ AQUI: Não processamos mais nada. Só avisamos o usuário.
        flash(f'Arquivo "{filename}" recebido! Ele será processado em breve e estará disponível para consulta.')
        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__':
    # Apenas carrega a chain existente quando o app inicia
    get_rag_chain()
    app.run(debug=True)