from tqdm import tqdm
import os
import chardet
import numpy as np
from pinecone import Pinecone
from langchain.docstore.document import Document as LangchainDocument
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import OpenAIEmbeddings
import spacy


# Constants
DATA_DIR = "data"
JOURNAL_DIR = "journals"

# Load environment variables
OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
PINECONE_API_KEY = os.environ["PINECONE_API_KEY"]
PINECONE_INDEX = os.environ["PINECONE_INDEX"]

# Initialize Pinecone
pc = Pinecone(api_key=PINECONE_API_KEY)
index = pc.Index(PINECONE_INDEX)

# Initialize embedding model
embedding_model = OpenAIEmbeddings(
    model="text-embedding-3-small",
    api_key=OPENAI_API_KEY
)

nlp = spacy.load("xx_sent_ud_sm")  # For multilingual support including Arabic

def sentence_overlap_chunks(text, chunk_size=2000):
    doc = nlp(text)
    sentences = [sent.text.strip() for sent in doc.sents]
    
    chunks = []
    i = 0

    while i < len(sentences):
        chunk = []
        length = 0
        start_i = i

        # Fill chunk up to chunk_size characters
        while i < len(sentences) and length + len(sentences[i]) <= chunk_size:
            chunk.append(sentences[i])
            length += len(sentences[i]) + 1
            i += 1

        # Join and store chunk
        if chunk:
            try:
                chunk = " ".join(chunk)
                chunks.append(chunk)
                i+=1
            except:
                print("can't process line.")
                i+=4

        # Overlap: start next chunk with last sentence of current chunk
        i = i - 3

    return chunks


# for data
for filename in os.listdir(DATA_DIR):
    if filename.endswith(".pdf"):
        filepath = os.path.join(DATA_DIR, filename)
        namespace = "ns4"

        print(f"Processing {filename} → namespace: {namespace}")

        reader = fitz.open(filepath)
        content = ""
        for page in reader:
            text = page.get_text()
            if text:
                content += text + "\n"

        # Wrap in Langchain doc
        docs_processed = sentence_overlap_chunks(content)
        
        # Embed and prepare for upsert
        upsert_data = []
        for i, chunk in tqdm(enumerate(docs_processed), total=len(docs_processed), desc="Embedding chunks"):
            vector = embedding_model.embed_query(chunk)
            upsert_data.append({
                "id": f"{filename[:-4]}_chunk_{i}",
                "values": vector,
                "metadata": {
                    "text": chunk,
                    "source": filename
                }
            })

        # Upsert to Pinecone under this file's namespace
        print(f"⬆️ Upserting {len(upsert_data)} vectors to namespace '{namespace}'...")
        index.upsert(vectors=upsert_data, namespace=namespace)
        print(f"✅ Done with {filename}\n")

# for journals
for filename in os.listdir(JOURNAL_DIR):
    if filename.endswith(".txt"):
        filepath = os.path.join(JOURNAL_DIR, filename)
        namespace = filename[:-4]

        print(f"Processing {filename} → namespace: {namespace}")

        # Detect encoding
        with open(filepath, "rb") as f:
            raw_data = f.read()
            encoding = chardet.detect(raw_data)['encoding']

        # Read file
        with open(filepath, "r", encoding=encoding) as f:
            content = f.read()

        # Wrap in Langchain doc
        docs_processed = sentence_overlap_chunks(content, chunk_size=400)
        
        # Embed and prepare for upsert
        upsert_data = []
        for i, chunk in tqdm(enumerate(docs_processed), total=len(docs_processed), desc="Embedding chunks"):
            vector = embedding_model.embed_query(chunk)
            upsert_data.append({
                "id": f"{namespace}_chunk_{i}",
                "values": vector,
                "metadata": {
                    "text": chunk,
                    "source": filename
                }
            })

        # Upsert to Pinecone under this file's namespace
        print(f"⬆️ Upserting {len(upsert_data)} vectors to namespace '{namespace}'...")
        index.upsert(vectors=upsert_data, namespace=namespace)
        print(f"✅ Done with {filename}\n")