import gradio as gr
import os
from utils.embeddings import EmbeddingModel
from utils.vector_store import VectorStore
from utils.rag_chain import RAGChain
from config import CHUNK_SIZE, CHUNK_OVERLAP

# Initialize components
embedding_model = EmbeddingModel()
vector_store = VectorStore()
vector_store.create_collection()

def load_and_process_data(file_path):
    """Load và xử lý dataset"""
    with open(file_path, 'r', encoding='utf-8') as f:
        text = f.read()
    
    # Chia thành chunks
    chunks = []
    for i in range(0, len(text), CHUNK_SIZE - CHUNK_OVERLAP):
        chunk = text[i:i + CHUNK_SIZE]
        chunks.append(chunk)
    
    # Tạo embeddings
    embeddings = embedding_model.embed_documents(chunks)
    
    # Lưu vào vector store
    vector_store.add_documents(chunks, embeddings)
    
    return len(chunks)

# Load data khi khởi động
if os.path.exists("data/your_dataset.txt"):
    num_chunks = load_and_process_data("data/your_dataset.txt")
    print(f"Đã load {num_chunks} chunks")

# Initialize RAG chain
rag_chain = RAGChain(vector_store, embedding_model)

def chatbot_response(message, history):
    """Xử lý tin nhắn và trả về response"""
    try:
        response = rag_chain.get_answer(message)
        return response
    except Exception as e:
        return f"Lỗi: {str(e)}"

# Tạo Gradio interface
demo = gr.ChatInterface(
    fn=chatbot_response,
    title="RAG Chatbot với Gemini",
    description="Chatbot sử dụng RAG (Retrieval-Augmented Generation) với Gemini API",
    examples=[
        "Xin chào!",
        "Hãy giải thích về RAG",
        "Thông tin trong dataset là gì?"
    ],
    theme="soft"
)

if __name__ == "__main__":
    demo.launch()