import os
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from dotenv import load_dotenv
from langfuse import observe, propagate_attributes
from langchain_qdrant import QdrantVectorStore
from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams

load_dotenv()

qdrant_url = os.environ.get("QDRANT_CLUSTER_ENDPOINT")
qdrant_api_key = os.environ.get("QDRANT_API_KEY")

def extract_pdfs_from_folder(folder_path):
    pdf_files = []
    for file_name in os.listdir(folder_path):
        if file_name.endswith(".pdf"):
            pdf_files.append(os.path.join(folder_path, file_name))
    
    extracted_texts = []
    for pdf_file in pdf_files:
        loader = PyPDFLoader(pdf_file)
        pages = loader.load()
        extracted_texts += pages
    
    return extracted_texts

@observe()
class LoadDoc():
    def __init__(self):
        self.qdrant_url = qdrant_url
        self.qdrant_api_key = qdrant_api_key
        self.qdrant_collection_name = "pac"
        self.qdrant_client = QdrantClient(
            url=self.qdrant_url,
            api_key=self.qdrant_api_key
        )
        self.data_path = "./data/"
        self.model_name = "BAAI/bge-large-en"
        self.model_kwargs = {'device': 'cpu'}
        self.encode_kwargs = {'normalize_embeddings': False}

    def load_data_into_quadrant(self):
        if not self.qdrant_client.collection_exists(self.qdrant_collection_name):
            self.qdrant_client.create_collection(
                collection_name=self.qdrant_collection_name,
                vectors_config=VectorParams(size=1024, distance=Distance.COSINE),
            )
        count_req = self.qdrant_client.count(
            collection_name=self.qdrant_collection_name,
            exact=True,
        )
        if count_req.count == 0:
            # step 1
            data = extract_pdfs_from_folder(self.data_path)
            text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=50)
            texts = text_splitter.split_documents(data)

            # step 2
            embeddings = HuggingFaceEmbeddings(
                model_name=self.model_name,
                model_kwargs=self.model_kwargs,
                encode_kwargs=self.encode_kwargs
            )

            # step 3
            qdrant = QdrantVectorStore.from_documents(
                texts,
                embeddings,
                url=self.qdrant_url,
                prefer_grpc=True,
                api_key=self.qdrant_api_key,
                collection_name=self.qdrant_collection_name,
                force_recreate=True
            )