import chromadb
from chromadb.utils import embedding_functions
import os

class BaitakMemory:
    def __init__(self, collection_name="baitak_insights", path="/home/ubuntu/Project-Baitak-Intelligence/data/chroma_db"):
        self.client = chromadb.PersistentClient(path=path)
        # Using OpenAI embeddings for simplicity, can be replaced with Gemini if API key is available

        self.collection = self.client.get_or_create_collection(name=collection_name)

    def add_insight(self, insight_text: str, metadata: dict = None):
        # ChromaDB requires a unique ID for each document
        doc_id = f"insight_{self.collection.count() + 1}"
        self.collection.add(
            documents=[insight_text],
            metadatas=[metadata if metadata else {}],
            ids=[doc_id]
        )
        print(f"Added insight with ID: {doc_id}")

    def query_insights(self, query_text: str, n_results: int = 5):
        results = self.collection.query(
            query_texts=[query_text],
            n_results=n_results
        )
        return results

    def get_all_insights(self):
        return self.collection.get(include=['documents', 'metadatas'])

    def clear_memory(self):
        self.client.delete_collection(name=self.collection.name)
        self.collection = self.client.get_or_create_collection(name=self.collection.name, embedding_function=self.embedding_function)
        print(f"Memory collection \'{self.collection.name}\' cleared.")

if __name__ == "__main__":
    memory = BaitakMemory()
    memory.clear_memory() # Clear for fresh start

    # Add some synthetic insights
    memory.add_insight("KFH's Net Financing Margin (NFM) remained stable around 3.5% in FY2025 due to CBK rate cuts.", {"quarter": "Q4-2025", "kpi": "NFM"})
    memory.add_insight("Digital transformation initiatives led to a 10% increase in mobile banking users in Q3-2025.", {"quarter": "Q3-2025", "kpi": "Digital Transformation"})
    memory.add_insight("The Cost-to-Income ratio improved to 34.06% in FY2025, down from 35.46% in FY2024.", {"quarter": "FY2025", "kpi": "Cost-to-Income"})

    # Query insights
    print("\nQuerying for NFM insights:")
    results = memory.query_insights("What was the Net Financing Margin in 2025?")
    for doc, meta in zip(results["documents"], results["metadatas"]):
        print(f"Insight: {doc}, Metadata: {meta}")

    print("\nQuerying for digital transformation:")
    results = memory.query_insights("Tell me about digital milestones.")
    for doc, meta in zip(results["documents"], results["metadatas"]):
        print(f"Insight: {doc}, Metadata: {meta}")