"""Episodic Memory Store — LanceDB vector storage with timestamp indexing.

Stores timestamped events and conversations as vector embeddings.
Supports semantic search with metadata filtering (confidence, time range).
Maps to the hippocampus in human cognitive architecture.
"""

from __future__ import annotations

import os
from datetime import datetime, timedelta, timezone
from typing import Optional

import numpy as np

from mnemosynth.core.types import MemoryNode, MemoryType, MemoryStatus
from mnemosynth.stores.base import BaseStore


class EpisodicStore(BaseStore):
    """LanceDB-backed vector store for episodic memories."""

    TABLE_NAME = "episodic_memories"

    def __init__(self, config, embeddings):
        self.config = config
        self.embeddings = embeddings
        self._db = None
        self._table = None

    @property
    def db(self):
        """Lazy-load LanceDB connection."""
        if self._db is None:
            import lancedb
            self._db = lancedb.connect(self.config.vectors_dir)
        return self._db

    @property
    def table(self):
        """Get or create the episodic memories table."""
        if self._table is None:
            try:
                self._table = self.db.open_table(self.TABLE_NAME)
            except Exception:
                # Table doesn't exist yet — create on first add
                self._table = None
        return self._table

    def _ensure_table(self, embedding_dim: int) -> None:
        """Create the table if it doesn't exist."""
        if self.table is None:
            import pyarrow as pa

            schema = pa.schema([
                pa.field("id", pa.string()),
                pa.field("content", pa.string()),
                pa.field("vector", pa.list_(pa.float32(), embedding_dim)),
                pa.field("confidence", pa.float64()),
                pa.field("created_at", pa.string()),
                pa.field("last_accessed", pa.string()),
                pa.field("access_count", pa.int64()),
                pa.field("status", pa.string()),
                pa.field("sentiment_score", pa.float64()),
                pa.field("corroboration_count", pa.int64()),
            ])
            self._table = self.db.create_table(self.TABLE_NAME, schema=schema)

    def add(self, memory: MemoryNode) -> None:
        """Store an episodic memory with its embedding vector."""
        if memory.embedding is None:
            memory.embedding = self.embeddings.embed(memory.content)

        self._ensure_table(len(memory.embedding))

        row = {
            "id": memory.id,
            "content": memory.content,
            "vector": memory.embedding,
            "confidence": memory.confidence,
            "created_at": memory.created_at.isoformat(),
            "last_accessed": memory.last_accessed.isoformat(),
            "access_count": memory.access_count,
            "status": memory.status.value,
            "sentiment_score": memory.sentiment_score,
            "corroboration_count": memory.corroboration_count,
        }
        self.table.add([row])

    def search(self, query: str, limit: int = 5) -> list[MemoryNode]:
        """Semantic search using vector similarity."""
        if self.table is None:
            return []

        query_embedding = self.embeddings.embed(query)

        try:
            results = (
                self.table.search(query_embedding)
                .where("status = 'active'", prefilter=True)
                .limit(limit)
                .to_pandas()
            )
        except Exception:
            # Fallback without filter if column not available
            try:
                results = (
                    self.table.search(query_embedding)
                    .limit(limit)
                    .to_pandas()
                )
            except Exception:
                return []

        memories = []
        for _, row in results.iterrows():
            node = MemoryNode(
                id=row["id"],
                content=row["content"],
                memory_type=MemoryType.EPISODIC,
                confidence=row.get("confidence", 0.85),
                created_at=datetime.fromisoformat(row["created_at"]),
                last_accessed=datetime.fromisoformat(row["last_accessed"]),
                access_count=int(row.get("access_count", 0)),
                status=MemoryStatus(row.get("status", "active")),
                sentiment_score=row.get("sentiment_score", 0.0),
                corroboration_count=int(row.get("corroboration_count", 0)),
                embedding=row["vector"] if "vector" in row else None,
            )
            memories.append(node)

        return memories

    def get(self, memory_id: str) -> Optional[MemoryNode]:
        """Retrieve a specific episodic memory by ID."""
        if self.table is None:
            return None

        try:
            results = self.table.search().where(f"id = '{memory_id}'").limit(1).to_pandas()
            if len(results) == 0:
                return None

            row = results.iloc[0]
            return MemoryNode(
                id=row["id"],
                content=row["content"],
                memory_type=MemoryType.EPISODIC,
                confidence=row.get("confidence", 0.85),
                created_at=datetime.fromisoformat(row["created_at"]),
                last_accessed=datetime.fromisoformat(row["last_accessed"]),
                access_count=int(row.get("access_count", 0)),
                status=MemoryStatus(row.get("status", "active")),
                embedding=row.get("vector"),
            )
        except Exception:
            return None

    def update(self, memory: MemoryNode) -> None:
        """Update an existing episodic memory."""
        # LanceDB doesn't support in-place update well, so delete + re-add
        self.delete(memory.id)
        self.add(memory)

    def delete(self, memory_id: str) -> bool:
        """Delete an episodic memory by ID."""
        if self.table is None:
            return False
        try:
            self.table.delete(f"id = '{memory_id}'")
            return True
        except Exception:
            return False

    def count(self) -> int:
        """Count total episodic memories."""
        if self.table is None:
            return 0
        try:
            return len(self.table)
        except Exception:
            return 0

    def get_recent(self, days: int = 7) -> list[MemoryNode]:
        """Get episodic memories from the last N days."""
        if self.table is None:
            return []

        try:
            # Get all and filter by date in Python
            df = self.table.to_pandas()
            cutoff = datetime.now(timezone.utc) - timedelta(days=days)

            memories = []
            for _, row in df.iterrows():
                created = datetime.fromisoformat(row["created_at"])
                if created >= cutoff and row.get("status", "active") == "active":
                    node = MemoryNode(
                        id=row["id"],
                        content=row["content"],
                        memory_type=MemoryType.EPISODIC,
                        confidence=row.get("confidence", 0.85),
                        created_at=created,
                        last_accessed=datetime.fromisoformat(row["last_accessed"]),
                        access_count=int(row.get("access_count", 0)),
                        status=MemoryStatus.ACTIVE,
                        embedding=row.get("vector"),
                    )
                    memories.append(node)

            return memories
        except Exception:
            return []
