"""
CodeContext Store - LanceDB operations with shared embedder support

Multi-project architecture:
- Server loads ONE embedding model at startup
- Each project gets its own ContextStore with separate database
- Stores share the same embedder (no model reload per project)

The ContextStore class uses mixins to organize functionality:
- ConflictMixin: Conflict detection and resolution
- ArchiveMixin: Archive and restore operations
- SummaryMixin: Summary version management
- ConsolidationMixin: Consolidation helper methods
"""

import json
from pathlib import Path
from typing import List, Dict, Optional

import lancedb
from .embedder import Embedder

from .gpu_utils import get_device, get_gpu_info
from .utils.timestamps import utc_now_iso, utc_timestamp_ms
from .store import (
    ConflictMixin,
    ArchiveMixin,
    SummaryMixin,
    ConsolidationMixin,
    parse_files
)
from .constants import EMBEDDING_BATCH_SIZE
from .logging_config import get_logger

logger = get_logger(__name__)


class ContextStore(ConflictMixin, ArchiveMixin, SummaryMixin, ConsolidationMixin):
    """
    Vector memory store with LanceDB - uses shared external embedder.

    Combines functionality from multiple mixins:
    - ConflictMixin: Conflict detection and resolution
    - ArchiveMixin: Archive, delete, restore operations
    - SummaryMixin: Summary version management
    - ConsolidationMixin: Consolidation helpers
    """

    VALID_TYPES = {"bugfix", "feature", "discovery", "decision", "refactor", "optimization", "gotcha", "resolution"}

    def __init__(
        self,
        db_path: str,
        embedder: Embedder,
        device: str,
        project_name: str = "default"
    ):
        """
        Initialize memory store for a specific project.

        Args:
            db_path: Path to this project's LanceDB storage
            embedder: Shared Embedder (loaded once by server)
            device: 'cuda' or 'cpu' (from server)
            project_name: Name of this project
        """
        self.db_path = Path(db_path)
        self.db_path.mkdir(parents=True, exist_ok=True)
        self.project_name = project_name
        self.device = device

        # Use shared embedder (not loaded here - passed from server)
        self.embedder = embedder
        self.embedding_dim = self.embedder.get_sentence_embedding_dimension()

        # Initialize LanceDB for this project
        self.db = lancedb.connect(str(self.db_path / "lancedb"))
        self.table_name = "memories"
        self.conflicts_table_name = "conflicts"
        self.summaries_table_name = "summaries"
        self.archive_table_name = "archive"
        self._init_table()
        self._init_conflicts_table()
        self._init_summaries_table()
        self._init_archive_table()
        logger.info(f"[Store:{project_name}] Ready (dim={self.embedding_dim})")

    def _init_table(self):
        """Initialize or open the memories table."""
        try:
            self.table = self.db.open_table(self.table_name)
            logger.info(f"[Store:{self.project_name}] Opened table: {self.count()} memories")
            # Check if migration needed for conflict fields
            self._migrate_schema_if_needed()
        except FileNotFoundError:
            # Table doesn't exist yet - this is expected for new projects
            self.table = None
            logger.info(f"[Store:{self.project_name}] New database, no memories yet")
        except Exception as e:
            # Unexpected error - log it but continue
            self.table = None
            logger.error(f"[Store:{self.project_name}] ERROR opening table: {e}")

    def _migrate_schema_if_needed(self):
        """
        Migrate from old embedded conflict schema to new separate conflicts table.

        Old memories had: conflict_status, conflict_with_id, conflict_resolution, conflict_detected_at
        New schema: conflicts go in separate 'conflicts' table, memories have no conflict fields.
        """
        if self.table is None:
            return

        try:
            df = self.table.to_pandas()
            if df.empty:
                return

            # Check if old conflict fields exist
            old_conflict_fields = ["conflict_status", "conflict_with_id", "conflict_resolution", "conflict_detected_at"]
            has_old_fields = any(field in df.columns for field in old_conflict_fields)

            if not has_old_fields:
                return  # Already migrated or fresh database

            # Migrate any pending conflicts to the new conflicts table
            pending_count = 0
            if "conflict_status" in df.columns:
                pending = df[df["conflict_status"] == "pending"]
                for _, r in pending.iterrows():
                    conflict_with_id = r.get("conflict_with_id", "")
                    if conflict_with_id:
                        # Get the existing memory data
                        existing_df = df[df["id"] == conflict_with_id]
                        if not existing_df.empty:
                            existing = existing_df.iloc[0]
                            # Create conflict record in new table
                            self._create_conflict(
                                new_memory={
                                    "id": r.get("id"),
                                    "type": r.get("type"),
                                    "title": r.get("title"),
                                    "fact": r.get("fact"),
                                    "files": r.get("files"),
                                    "created_at": r.get("created_at")
                                },
                                existing_memory={
                                    "id": existing.get("id"),
                                    "type": existing.get("type"),
                                    "title": existing.get("title"),
                                    "fact": existing.get("fact"),
                                    "files": existing.get("files"),
                                    "created_at": existing.get("created_at")
                                }
                            )
                            pending_count += 1

            # Remove old conflict columns from memories
            columns_to_keep = [c for c in df.columns if c not in old_conflict_fields]
            df = df[columns_to_keep]

            # Recreate the table without conflict fields
            records = df.to_dict('records')

            # Drop and recreate (LanceDB doesn't support ALTER TABLE)
            self.db.drop_table(self.table_name)
            self.table = self.db.create_table(self.table_name, records)

            logger.info(f"[Store:{self.project_name}] Migration complete: removed conflict fields, migrated {pending_count} pending conflicts")
        except Exception as e:
            logger.debug(f"[Store:{self.project_name}] Migration check failed (may be OK): {e}")

    def _init_conflicts_table(self):
        """Initialize or open the conflicts table."""
        # Required fields for the current schema (including ranking metadata)
        required_fields = {"id", "status", "distance", "file_overlap", "content_length_diff", "confidence"}

        try:
            self.conflicts_table = self.db.open_table(self.conflicts_table_name)

            # Check if table has required fields (schema migration)
            df = self.conflicts_table.to_pandas()
            if not df.empty:
                existing_fields = set(df.columns)
                missing_fields = required_fields - existing_fields

                if missing_fields:
                    logger.info(f"[Store:{self.project_name}] Conflicts table outdated (missing: {missing_fields})")
                    logger.info(f"[Store:{self.project_name}] Dropping old table and recreating on next conflict...")
                    self.db.drop_table(self.conflicts_table_name)
                    self.conflicts_table = None
                else:
                    logger.debug(f"[Store:{self.project_name}] Opened conflicts table")
            else:
                logger.debug(f"[Store:{self.project_name}] Opened conflicts table (empty)")
        except Exception:
            self.conflicts_table = None  # Will be created on first conflict
            logger.debug(f"[Store:{self.project_name}] Conflicts table: will create on first conflict")

    def _init_summaries_table(self):
        """Initialize or open the summaries table for version history."""
        try:
            self.summaries_table = self.db.open_table(self.summaries_table_name)
            count = len(self.summaries_table)
            logger.debug(f"[Store:{self.project_name}] Opened summaries table ({count} versions)")
        except Exception:
            self.summaries_table = None  # Will be created on first summary
            logger.debug(f"[Store:{self.project_name}] Summaries table: will create on first summary")

    def _init_archive_table(self):
        """Initialize or open the archive table for deleted memories."""
        try:
            self.archive_table = self.db.open_table(self.archive_table_name)
            count = len(self.archive_table)
            logger.debug(f"[Store:{self.project_name}] Opened archive table ({count} entries)")
        except Exception:
            self.archive_table = None  # Created on first archive
            logger.debug(f"[Store:{self.project_name}] Archive table: will create on first archive")

    def add(
        self,
        obs_type: str,
        title: str,
        fact: str,
        files: List[str],
        created_at: str = None
    ) -> Dict:
        """
        Add a single memory to this project's database.

        Returns:
            Dict with keys:
            - id: Memory ID (new or existing if duplicate)
            - status: "added" | "duplicate" | "conflict_flagged"
            - existing_match: Dict of matching memory (if duplicate/conflict)
            - distance: L2 distance to nearest neighbor
        """
        if obs_type not in self.VALID_TYPES:
            obs_type = "discovery"

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

        # Check for duplicates/conflicts
        similarity = self._check_similarity(embedding)

        if similarity["status"] == "duplicate":
            return {
                "id": similarity["match"]["id"],
                "status": "duplicate",
                "existing_match": similarity["match"],
                "distance": similarity["distance"]
            }

        # Prepare memory record (no conflict fields - conflicts go in separate table)
        is_conflict = similarity["status"] == "conflict"
        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])

        result = {
            "id": record["id"],
            "status": "conflict_flagged" if is_conflict else "added",
            "distance": similarity["distance"]
        }

        # If conflict detected, create conflict record in separate table
        if is_conflict:
            conflict_id = self._create_conflict(
                new_memory=record,
                existing_memory=similarity["match"],
                distance=similarity["distance"]
            )
            result["existing_match"] = similarity["match"]
            result["conflict_id"] = conflict_id

        return result

    def add_batch(self, memories: List[Dict]) -> Dict:
        """
        Add multiple memories at once (more efficient).

        Args:
            memories: List of dicts with type, title, fact, files

        Returns:
            Dict with keys:
            - added: Number of memories added
            - duplicates: List of duplicate matches
            - conflicts: List of conflict-flagged memories
        """
        if not memories:
            return {"added": 0, "duplicates": [], "conflicts": []}

        # Prepare texts for batch encoding
        texts = []
        for m in memories:
            obs_type = m.get("type", "discovery")
            if obs_type not in self.VALID_TYPES:
                obs_type = "discovery"
            text = f"[{obs_type}] {m['title']}\n{m['fact']}"
            texts.append(text)

        # Batch encode (much faster than one at a time)
        embeddings = self.embedder.encode(texts, batch_size=EMBEDDING_BATCH_SIZE, show_progress_bar=False)

        # Check each memory for duplicates/conflicts
        records_to_add = []
        duplicates = []
        conflicts = []
        timestamp = utc_timestamp_ms()

        for i, (m, text, emb) in enumerate(zip(memories, texts, embeddings)):
            emb_list = emb.tolist()
            similarity = self._check_similarity(emb_list)

            obs_type = m.get("type", "discovery")
            if obs_type not in self.VALID_TYPES:
                obs_type = "discovery"

            if similarity["status"] == "duplicate":
                duplicates.append({
                    "input": {"title": m["title"], "fact": m["fact"]},
                    "existing_match": similarity["match"],
                    "distance": similarity["distance"]
                })
                continue  # Skip adding duplicate

            # Determine if conflict
            is_conflict = similarity["status"] == "conflict"
            now = utc_now_iso()

            # Memory record (no conflict fields - conflicts go in separate table)
            record = {
                "id": f"{self.project_name}_{timestamp + i}",
                "type": obs_type,
                "title": m["title"],
                "fact": m["fact"],
                "files": json.dumps(m.get("files", [])),
                "project": self.project_name,
                "text": text,
                "created_at": now,
                "vector": emb_list
            }

            records_to_add.append(record)

            if is_conflict:
                # Create conflict record in separate table
                conflict_id = self._create_conflict(
                    new_memory=record,
                    existing_memory=similarity["match"],
                    distance=similarity["distance"]
                )
                conflicts.append({
                    "id": record["id"],
                    "conflict_id": conflict_id,
                    "title": m["title"],
                    "fact": m["fact"],
                    "existing_match": similarity["match"],
                    "distance": similarity["distance"]
                })

        # Add non-duplicate records
        if records_to_add:
            if self.table is None:
                self.table = self.db.create_table(self.table_name, records_to_add)
            else:
                self.table.add(records_to_add)

        return {
            "added": len(records_to_add),
            "duplicates": duplicates,
            "conflicts": conflicts
        }

    def query(self, query: str, top_k: int = 5) -> List[Dict]:
        """
        Query memories by semantic similarity.

        Args:
            query: Search query
            top_k: Number of results

        Returns:
            List of matching memories with relevance scores
        """
        if self.table is None or self.count() == 0:
            return []

        embedding = self.embedder.encode(query).tolist()
        results = self.table.search(embedding).limit(top_k).to_list()

        return [
            {
                "id": r.get("id"),
                "type": r["type"],
                "title": r["title"],
                "fact": r["fact"],
                "files": parse_files(r.get("files")),
                "project": self.project_name,
                "created_at": r.get("created_at"),
                "distance": r.get("_distance", 0),
                "relevance_pct": max(0, min(100, 100 - (r.get("_distance", 0) * 40)))
            }
            for r in results
        ]

    def query_filtered(
        self,
        query: str,
        top_k: int = 5,
        type_filter: Optional[List[str]] = None,
        file_filter: Optional[str] = None,
        sort_order: str = "relevance"
    ) -> List[Dict]:
        """
        Query memories by semantic similarity with optional filtering and sorting.

        Args:
            query: Search query
            top_k: Number of results to return
            type_filter: Optional list of types to include (e.g., ["bugfix", "feature"])
            file_filter: Optional partial match string for file paths (case-insensitive)
            sort_order: "relevance" (default), "newest", or "oldest"

        Returns:
            List of matching memories with relevance scores
        """
        if self.table is None or self.count() == 0:
            return []

        # Fetch more results if filtering to account for filtered-out items
        fetch_multiplier = 3 if type_filter or file_filter else 1
        fetch_limit = min(top_k * fetch_multiplier, self.count())

        embedding = self.embedder.encode(query).tolist()
        results = self.table.search(embedding).limit(fetch_limit).to_list()

        # Format results
        formatted = [
            {
                "id": r.get("id"),
                "type": r["type"],
                "title": r["title"],
                "fact": r["fact"],
                "files": parse_files(r.get("files")),
                "project": self.project_name,
                "created_at": r.get("created_at"),
                "distance": r.get("_distance", 0),
                "relevance_pct": max(0, min(100, 100 - (r.get("_distance", 0) * 40)))
            }
            for r in results
        ]

        # Filter by type if specified (case-insensitive)
        if type_filter:
            normalized_filter = {t.lower() for t in type_filter}
            formatted = [r for r in formatted if r.get("type", "").lower() in normalized_filter]

        # Filter by file if specified (partial match, case-insensitive)
        if file_filter:
            pattern = file_filter.lower()
            formatted = [
                r for r in formatted
                if any(pattern in f.lower() for f in r.get("files", []))
            ]

        # Sort based on sort_order
        if sort_order == "newest":
            formatted.sort(key=lambda x: x.get("created_at", ""), reverse=True)
        elif sort_order == "oldest":
            formatted.sort(key=lambda x: x.get("created_at", ""))
        # "relevance" keeps original distance-based order

        # Limit to top_k after filtering
        return formatted[:top_k]

    def batch_query(self, queries: List[str], top_k: int = 5) -> List[Dict]:
        """
        Query multiple searches at once (more efficient).

        Args:
            queries: List of search queries
            top_k: Results per query

        Returns:
            List of result sets, one per query
        """
        if self.table is None or self.count() == 0:
            return [{"query": q, "results": []} for q in queries]

        # Batch encode queries
        embeddings = self.embedder.encode(queries, batch_size=EMBEDDING_BATCH_SIZE, show_progress_bar=False)

        # Execute each search
        results = []
        for query, emb in zip(queries, embeddings):
            search_results = self.table.search(emb.tolist()).limit(top_k).to_list()
            formatted = [
                {
                    "id": r.get("id"),
                    "type": r["type"],
                    "title": r["title"],
                    "fact": r["fact"],
                    "files": parse_files(r.get("files")),
                    "project": self.project_name,
                    "distance": r.get("_distance", 0),
                    "relevance_pct": max(0, min(100, 100 - (r.get("_distance", 0) * 40)))
                }
                for r in search_results
            ]
            results.append({"query": query, "results": formatted})

        return results

    def list_all(
        self,
        limit: int = 20,
        offset: int = 0,
        type_filter: Optional[List[str]] = None,
        file_filter: Optional[str] = None,
        sort_order: str = "newest"
    ) -> List[Dict]:
        """List all memories (for dashboard) with optional filtering and sorting.

        Args:
            limit: Max number of memories to return
            offset: Number of memories to skip
            type_filter: Optional list of types to include (e.g., ["bugfix", "feature"])
            file_filter: Optional partial match string for file paths (case-insensitive)
            sort_order: "newest" (default) or "oldest"
        """
        if self.table is None:
            return []

        try:
            results = self.table.to_pandas().to_dict('records')

            # Filter by type if specified
            if type_filter:
                results = [r for r in results if r.get("type") in type_filter]

            # Filter by file if specified (partial match, case-insensitive)
            if file_filter:
                pattern = file_filter.lower()
                results = [
                    r for r in results
                    if any(pattern in f.lower() for f in parse_files(r.get("files")))
                ]

            # Sort by created_at
            reverse = sort_order != "oldest"
            results.sort(key=lambda x: x.get("created_at", ""), reverse=reverse)

            # Limit and format
            return [
                {
                    "id": r.get("id"),
                    "type": r["type"],
                    "title": r["title"],
                    "fact": r["fact"],
                    "files": parse_files(r.get("files")),
                    "project": r.get("project", self.project_name),
                    "created_at": r.get("created_at"),
                }
                for r in results[offset:offset + limit]
            ]
        except Exception as e:
            logger.error(f"[Store:{self.project_name}] Error listing: {e}")
            return []

    def count(self) -> int:
        """Return number of memories."""
        if self.table is None:
            return 0
        return len(self.table)

    def get_unique_files(self) -> List[str]:
        """Get all unique file paths across all memories for autocomplete."""
        if self.table is None:
            return []

        try:
            results = self.table.to_pandas().to_dict('records')
            files_set = set()
            for r in results:
                files_set.update(parse_files(r.get("files")))
            return sorted(list(files_set))
        except Exception as e:
            logger.error(f"[Store:{self.project_name}] Error getting unique files: {e}")
            return []

    def get_info(self) -> Dict:
        """Get store information."""
        gpu_info = get_gpu_info()
        return {
            "project": self.project_name,
            "db_path": str(self.db_path),
            "device": self.device,
            "embedding_dim": self.embedding_dim,
            "memory_count": self.count(),
            "pending_conflicts": self.count_pending_conflicts(),
            "gpu": gpu_info
        }

    def get_memory_by_id(self, memory_id: str) -> Optional[Dict]:
        """Get a single memory by ID."""
        if self.table is None:
            return None

        try:
            df = self.table.to_pandas()
            match = df[df["id"] == memory_id]
            if match.empty:
                return None

            r = match.iloc[0]
            return {
                "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)
            }
        except Exception as e:
            logger.error(f"[Store:{self.project_name}] Error getting memory: {e}")
            return None
