"""
Consolidation Operations Mixin for ContextStore.

Provides helper methods for memory consolidation:
- get_all_memories_with_vectors: Get all memories including embeddings
- find_similar_memories: Find semantically similar memories
- delete_memories: Bulk delete by IDs
- add_consolidated_memory: Add a merged memory
"""

import json
from typing import Dict, List

from ..utils.timestamps import utc_now_iso, utc_timestamp_ms
from .utils import parse_files


class ConsolidationMixin:
    """
    Mixin class providing consolidation helper operations.

    Requires the base class to have:
    - self.table: LanceDB memories table (or None)
    - self.embedder: Embedder instance
    - self.project_name: str
    - self.VALID_TYPES: set of valid memory types
    """

    def get_all_memories_with_vectors(self) -> List[Dict]:
        """
        Get all memories including their vector embeddings.

        Used by consolidation to build similarity clusters.

        Returns:
            List of memory dicts with 'vector' field included
        """
        if self.table is None:
            return []

        try:
            df = self.table.to_pandas()
            results = []
            for _, r in df.iterrows():
                results.append({
                    "id": r.get("id"),
                    "type": r.get("type"),
                    "title": r.get("title"),
                    "fact": r.get("fact"),
                    "files": parse_files(r.get("files")),
                    "created_at": r.get("created_at"),
                    "project": r.get("project", self.project_name),
                    "vector": r.get("vector")  # Include embedding
                })
            return results
        except Exception as e:
            print(f"[Store:{self.project_name}] Error getting memories with vectors: {e}")
            return []

    def find_similar_memories(
        self,
        memory_id: str,
        threshold: float = 0.55,
        limit: int = 50
    ) -> List[Dict]:
        """
        Find memories similar to a given memory.

        Uses LanceDB vector search to find semantically similar entries.

        Args:
            memory_id: ID of the memory to find similar ones for
            threshold: Maximum distance (L2) to consider similar (default 0.55)
            limit: Maximum number of similar memories to return

        Returns:
            List of similar memories with distance scores
        """
        if self.table is None:
            return []

        # Get the source memory's vector
        try:
            df = self.table.to_pandas()
            match = df[df["id"] == memory_id]
            if match.empty:
                return []

            source_vector = match.iloc[0].get("vector")
            if source_vector is None:
                return []

            # Search for similar memories
            results = self.table.search(source_vector).limit(limit + 1).to_list()

            # Filter by threshold and exclude self
            similar = []
            for r in results:
                if r.get("id") == memory_id:
                    continue  # Skip self
                distance = r.get("_distance", float('inf'))
                if distance <= threshold:
                    similar.append({
                        "id": r.get("id"),
                        "type": r.get("type"),
                        "title": r.get("title"),
                        "fact": r.get("fact"),
                        "files": parse_files(r.get("files")),
                        "created_at": r.get("created_at"),
                        "distance": distance
                    })

            return similar
        except Exception as e:
            print(f"[Store:{self.project_name}] Error finding similar memories: {e}")
            return []

    def delete_memories(self, ids: List[str]) -> int:
        """
        Delete multiple memories by their IDs.

        Args:
            ids: List of memory IDs to delete

        Returns:
            Number of memories successfully deleted
        """
        if self.table is None or not ids:
            return 0

        deleted = 0
        for memory_id in ids:
            try:
                self.table.delete(f'id = "{memory_id}"')
                deleted += 1
            except Exception as e:
                print(f"[Store:{self.project_name}] Error deleting {memory_id}: {e}")

        return deleted

    def add_consolidated_memory(
        self,
        obs_type: str,
        title: str,
        fact: str,
        files: List[str],
        source_ids: List[str],
        created_at: str = None
    ) -> Dict:
        """
        Add a consolidated memory (merged from multiple sources).

        Similar to add() but:
        - Skips duplicate/conflict detection (we're replacing duplicates)
        - Records source IDs for audit trail

        Args:
            obs_type: Memory type
            title: Consolidated title
            fact: Consolidated fact
            files: Combined files list
            source_ids: IDs of memories that were consolidated
            created_at: Optional timestamp (defaults to now)

        Returns:
            Dict with new memory ID
        """
        if obs_type not in self.VALID_TYPES:
            obs_type = "discovery"

        text = f"[{obs_type}] {title}\n{fact}"
        embedding = self.embedder.encode(text).tolist()

        now = created_at if created_at else utc_now_iso()
        record = {
            "id": f"{self.project_name}_{utc_timestamp_ms()}",
            "type": obs_type,
            "title": title,
            "fact": fact,
            "files": json.dumps(files),
            "project": self.project_name,
            "text": text,
            "created_at": now,
            "vector": embedding
        }

        if self.table is None:
            self.table = self.db.create_table(self.table_name, [record])
        else:
            self.table.add([record])

        return {
            "id": record["id"],
            "source_ids": source_ids,
            "status": "consolidated"
        }
