# src/utils.py
import os
import json
from src.memory import MemoryManager  # Corrected import path
from src.llm_interface import LLMInterface  # Import LLMInterface
import logging
import spacy
from sklearn.cluster import KMeans
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def chunk_text(text, chunk_size=1000, overlap=100):
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start = end - overlap
    return chunks

def extract_and_summarize(query: str, memory_manager: MemoryManager, llm_interface: LLMInterface, system_prompt: str = "", max_tokens: int = 512, temperature: float = 0.7, top_p: float = 0.95) -> str:
    # Retrieve relevant memories from the database
    relevant_memories = memory_manager.retrieve_relevant_memories(query, limit=30)
    logging.info(f"Retrieved {len(relevant_memories)} relevant memories for query: {query}")

    # Combine relevant memories into a single context
    context = " ".join([memory['description'] for memory in relevant_memories])
    logging.info(f"Built context: {context}")

    # Truncate the context if it exceeds the token limit
    max_context_length = 30000  # Adjust this based on your LLM's token limit
    if len(context) > max_context_length:
        context = context[:max_context_length]
        logging.info(f"Truncated context to {max_context_length} characters.")

    # Use spaCy to generate sentence embeddings
    nlp = spacy.load('en_core_web_lg')
    sentences = context.split('.')
    sentence_embeddings = [nlp(sent).vector for sent in sentences]

    # Cluster sentences
    num_clusters = min(len(sentences), 10)  # Adjust the number of clusters
    kmeans = KMeans(n_clusters=num_clusters)
    kmeans.fit(sentence_embeddings)
    labels = kmeans.labels_

    # Select representative sentences from each cluster
    representative_sentences = []
    for i in range(num_clusters):
        cluster_sentences = [sentences[j] for j in range(len(sentences)) if labels[j] == i]
        if cluster_sentences:
            representative_sentences.append(max(cluster_sentences, key=len))  # Select the longest sentence as representative

    # Combine representative sentences to form a summary
    summary = " ".join(representative_sentences)
    logging.info(f"Generated summary: {summary}")

    # Use LLM to refine the summary
    try:
        refined_summary = llm_interface.send_message(f"Context: {summary}\nQuestion: {query}", system_prompt=system_prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p)
        logging.info(f"Refined summary: {refined_summary}")
    except Exception as e:
        refined_summary = f"Error refining summary: {e}"
        logging.error(f"Error refining summary: {e}")

    return refined_summary