import os
import json
import chromadb
import uuid

from sentence_transformers import SentenceTransformer

model = SentenceTransformer(
    "Snowflake/snowflake-arctic-embed-m-long",
    device="cpu",               # <-- the only line that changes
    trust_remote_code=True
)


import math
def embed_text(text, specific_context, max_tokens=2048, overlap=256):
    """
    Embeds the given text using the SentenceTransformer model.
    Splits the text into chunks if it exceeds max_tokens.
    """
    if len(text) + len(specific_context) <= max_tokens:
        return (text + specific_context, model.encode([text, specific_context], normalize_embeddings=True))

    # Split text into chunks
    chunks = []
    start = 0

    chunkNum = math.ceil(len(text) / (max_tokens - len(specific_context)))
    # 1000 + 20 = 1020
    # 3000 + 20 = 3020 -> 3000 / (2048 - 20) = 1.47 ~ 2

    for i in range(chunkNum):
        if start >= len(text):
            break
        # Calculate end index for the chunk
        end = min(start + max_tokens - len(specific_context), len(text))
        chunk = text[start:end] + specific_context
        chunks.append(chunk)
        start += (max_tokens - overlap)  # Overlap for next chunk

    return (chunks, model.encode(chunks, normalize_embeddings=True))


client = chromadb.HttpClient(host="localhost", port=8000)

# client.delete_collection("news")  # Replace "news" with your collection name

collection = client.get_or_create_collection("news")

# Load all the processed articles from the output folder
output_folder = os.path.join(os.path.dirname(__file__), "../ai_processor/output")
output_articles = []
# When embedding and upserting, use the source from the processed result
for filename in os.listdir(output_folder):
    if filename.endswith(".json"):
        with open(os.path.join(output_folder, filename), 'r', encoding='utf-8') as f:
            prev_proc = json.load(f)
            output_articles.append(prev_proc)

            def makeTickerMovement(ticker, movement, percent_change):
                """
                Helper function to create a ticker movement entry.
                """
                ticker = "" if ticker is None else ticker
                movement = "" if movement is None else movement
                percent_change = "" if percent_change is None else percent_change
                return [ticker, movement, percent_change]

            # Embed the article content
            try:
                specific_context = prev_proc['summary'] + ",".join([' '.join(makeTickerMovement(*x)) for x in prev_proc['tickersAndMovements']])
                chunks, embedded_content = embed_text(prev_proc['original_content'], prev_proc['summary'])
                
                print(f"Embedded content for {filename}: \n {embedded_content[:10]}...")  # Print first 10 values for preview

                collection.upsert(
                    ids      = [ str(uuid.uuid4()) for d in chunks ],  # Generate unique IDs for each embedded content
                    documents = chunks,                               # optional but nice for debugging
                    embeddings = embedded_content.tolist(),           # Chroma expects List[List[float]]
                    metadatas =[{
                        "source": prev_proc.get('source', 'Unknown'),  # Use the source from the processed result
                        "published": prev_proc.get('published', 'Unknown'),
                        "filename": prev_proc.get('filename', 'Unknown')  # Add filename for traceability
                    } for d in chunks],  # Metadata for each embedded content
                )
                
            except Exception as e:
                print(f"Error embedding content for {filename}: {e}")
                continue