"""Vector memory store using SQLite with vector support.

Provides persistent storage and similarity search for memory chunks.
"""

from __future__ import annotations

import json
import math
import sqlite3
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any


@dataclass
class MemoryChunk:
    """A chunk of memory content with metadata."""

    content: str
    source: str  # e.g., "MEMORY.md", "conversation", "file"
    chunk_type: str  # e.g., "principle", "fact", "event", "note"
    created_at: datetime = field(default_factory=datetime.now)
    metadata: dict[str, Any] = field(default_factory=dict)
    id: int | None = None

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "id": self.id,
            "content": self.content,
            "source": self.source,
            "chunk_type": self.chunk_type,
            "created_at": self.created_at.isoformat(),
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> MemoryChunk:
        """Create from dictionary."""
        return cls(
            id=data.get("id"),
            content=data["content"],
            source=data["source"],
            chunk_type=data["chunk_type"],
            created_at=datetime.fromisoformat(data["created_at"]),
            metadata=data.get("metadata", {}),
        )


@dataclass
class SearchResult:
    """Result from a similarity search."""

    chunk: MemoryChunk
    distance: float
    similarity: float

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "chunk": self.chunk.to_dict(),
            "distance": self.distance,
            "similarity": self.similarity,
        }


class VectorMemoryStore:
    """Vector memory store using SQLite.

    Provides persistent storage and similarity search for memory chunks.
    Uses Python-based cosine similarity for search.
    """

    def __init__(
        self,
        db_path: Path | None = None,
        embedding_dimension: int = 2048,
    ) -> None:
        """Initialize the vector memory store.

        Args:
            db_path: Path to SQLite database. Defaults to ~/.openclaw/oclawma/memory.db
            embedding_dimension: Dimension of embedding vectors.
        """
        if db_path is None:
            home = Path.home()
            db_path = home / ".openclaw" / "oclawma" / "memory.db"

        self.db_path = db_path
        self.embedding_dimension = embedding_dimension

        # Ensure directory exists
        self.db_path.parent.mkdir(parents=True, exist_ok=True)

        # Initialize database
        self._init_db()

    def _init_db(self) -> None:
        """Initialize database schema."""
        with sqlite3.connect(self.db_path) as conn:
            # Create memory chunks table
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS memory_chunks (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    content TEXT NOT NULL,
                    source TEXT NOT NULL,
                    chunk_type TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    metadata TEXT DEFAULT '{}'
                )
            """
            )

            # Create index on source for faster filtering
            conn.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_chunks_source
                ON memory_chunks(source)
            """
            )

            conn.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_chunks_type
                ON memory_chunks(chunk_type)
            """
            )

            # Create table for vector embeddings (stored as JSON)
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS memory_embeddings (
                    chunk_id INTEGER PRIMARY KEY,
                    embedding TEXT NOT NULL,
                    FOREIGN KEY (chunk_id) REFERENCES memory_chunks(id) ON DELETE CASCADE
                )
            """
            )

    def add_chunk(
        self,
        chunk: MemoryChunk,
        embedding: list[float],
    ) -> int:
        """Add a memory chunk with its embedding.

        Args:
            chunk: Memory chunk to add.
            embedding: Vector embedding of the chunk.

        Returns:
            ID of the inserted chunk.
        """
        with sqlite3.connect(self.db_path) as conn:
            # Insert chunk
            cursor = conn.execute(
                """
                INSERT INTO memory_chunks (content, source, chunk_type, created_at, metadata)
                VALUES (?, ?, ?, ?, ?)
                """,
                (
                    chunk.content,
                    chunk.source,
                    chunk.chunk_type,
                    chunk.created_at.isoformat(),
                    json.dumps(chunk.metadata),
                ),
            )
            chunk_id = cursor.lastrowid

            # Insert embedding
            conn.execute(
                "INSERT INTO memory_embeddings (chunk_id, embedding) VALUES (?, ?)",
                (chunk_id, json.dumps(embedding)),
            )

            return chunk_id or 0

    def add_chunks(
        self,
        chunks: list[MemoryChunk],
        embeddings: list[list[float]],
    ) -> list[int]:
        """Add multiple memory chunks with their embeddings.

        Args:
            chunks: List of memory chunks.
            embeddings: List of embeddings corresponding to chunks.

        Returns:
            List of IDs for inserted chunks.
        """
        if len(chunks) != len(embeddings):
            raise ValueError("Number of chunks must match number of embeddings")

        ids = []
        with sqlite3.connect(self.db_path) as conn:
            for chunk, embedding in zip(chunks, embeddings):
                cursor = conn.execute(
                    """
                    INSERT INTO memory_chunks (content, source, chunk_type, created_at, metadata)
                    VALUES (?, ?, ?, ?, ?)
                    """,
                    (
                        chunk.content,
                        chunk.source,
                        chunk.chunk_type,
                        chunk.created_at.isoformat(),
                        json.dumps(chunk.metadata),
                    ),
                )
                chunk_id = cursor.lastrowid

                conn.execute(
                    "INSERT INTO memory_embeddings (chunk_id, embedding) VALUES (?, ?)",
                    (chunk_id, json.dumps(embedding)),
                )

                if chunk_id:
                    ids.append(chunk_id)

        return ids

    def _cosine_similarity(self, a: list[float], b: list[float]) -> float:
        """Calculate cosine similarity between two vectors."""
        dot_product = sum(x * y for x, y in zip(a, b))
        norm_a = math.sqrt(sum(x * x for x in a))
        norm_b = math.sqrt(sum(x * x for x in b))

        if norm_a == 0 or norm_b == 0:
            return 0.0

        return dot_product / (norm_a * norm_b)

    def search(
        self,
        query_embedding: list[float],
        top_k: int = 5,
        source_filter: str | None = None,
        chunk_type_filter: str | None = None,
        min_similarity: float = 0.0,
    ) -> list[SearchResult]:
        """Search for similar memory chunks using cosine similarity.

        Args:
            query_embedding: Query vector embedding.
            top_k: Number of results to return.
            source_filter: Optional filter by source.
            chunk_type_filter: Optional filter by chunk type.
            min_similarity: Minimum similarity threshold (0-1).

        Returns:
            List of search results sorted by similarity.
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row

            # Build query with optional filters
            query = """
                SELECT
                    c.id,
                    c.content,
                    c.source,
                    c.chunk_type,
                    c.created_at,
                    c.metadata,
                    e.embedding
                FROM memory_chunks c
                JOIN memory_embeddings e ON c.id = e.chunk_id
            """
            params: list[Any] = []

            where_clauses = []
            if source_filter:
                where_clauses.append("c.source = ?")
                params.append(source_filter)

            if chunk_type_filter:
                where_clauses.append("c.chunk_type = ?")
                params.append(chunk_type_filter)

            if where_clauses:
                query += " WHERE " + " AND ".join(where_clauses)

            cursor = conn.execute(query, params)
            rows = cursor.fetchall()

            # Calculate similarities
            results = []
            for row in rows:
                embedding = json.loads(row["embedding"])
                similarity = self._cosine_similarity(query_embedding, embedding)

                if similarity < min_similarity:
                    continue

                # Distance is inverse of similarity for compatibility
                distance = 1.0 - similarity

                chunk = MemoryChunk(
                    id=row["id"],
                    content=row["content"],
                    source=row["source"],
                    chunk_type=row["chunk_type"],
                    created_at=datetime.fromisoformat(row["created_at"]),
                    metadata=json.loads(row["metadata"] or "{}"),
                )

                results.append(
                    SearchResult(
                        chunk=chunk,
                        distance=distance,
                        similarity=similarity,
                    )
                )

            # Sort by similarity (descending) and return top_k
            results.sort(key=lambda r: r.similarity, reverse=True)
            return results[:top_k]

    def get_chunk(self, chunk_id: int) -> MemoryChunk | None:
        """Get a memory chunk by ID.

        Args:
            chunk_id: ID of the chunk.

        Returns:
            Memory chunk if found, None otherwise.
        """
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute("SELECT * FROM memory_chunks WHERE id = ?", (chunk_id,))
            row = cursor.fetchone()

            if row is None:
                return None

            return MemoryChunk(
                id=row["id"],
                content=row["content"],
                source=row["source"],
                chunk_type=row["chunk_type"],
                created_at=datetime.fromisoformat(row["created_at"]),
                metadata=json.loads(row["metadata"] or "{}"),
            )

    def delete_chunk(self, chunk_id: int) -> bool:
        """Delete a memory chunk.

        Args:
            chunk_id: ID of the chunk to delete.

        Returns:
            True if deleted, False if not found.
        """
        with sqlite3.connect(self.db_path) as conn:
            # Delete from chunks (embedding will be deleted via CASCADE)
            cursor = conn.execute("DELETE FROM memory_chunks WHERE id = ?", (chunk_id,))

            return cursor.rowcount > 0

    def delete_by_source(self, source: str) -> int:
        """Delete all chunks from a source.

        Args:
            source: Source to delete.

        Returns:
            Number of chunks deleted.
        """
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute("DELETE FROM memory_chunks WHERE source = ?", (source,))

            return cursor.rowcount

    def get_stats(self) -> dict[str, Any]:
        """Get statistics about the memory store.

        Returns:
            Dictionary with statistics.
        """
        with sqlite3.connect(self.db_path) as conn:
            # Total chunks
            cursor = conn.execute("SELECT COUNT(*) FROM memory_chunks")
            total_chunks = cursor.fetchone()[0]

            # Chunks by source
            cursor = conn.execute("SELECT source, COUNT(*) FROM memory_chunks GROUP BY source")
            by_source = {row[0]: row[1] for row in cursor.fetchall()}

            # Chunks by type
            cursor = conn.execute(
                "SELECT chunk_type, COUNT(*) FROM memory_chunks GROUP BY chunk_type"
            )
            by_type = {row[0]: row[1] for row in cursor.fetchall()}

            # Date range
            cursor = conn.execute("SELECT MIN(created_at), MAX(created_at) FROM memory_chunks")
            row = cursor.fetchone()
            date_range = (
                {
                    "oldest": row[0],
                    "newest": row[1],
                }
                if row and row[0]
                else None
            )

            return {
                "total_chunks": total_chunks,
                "by_source": by_source,
                "by_type": by_type,
                "date_range": date_range,
                "db_path": str(self.db_path),
                "embedding_dimension": self.embedding_dimension,
            }

    def clear_all(self) -> None:
        """Clear all memory chunks. Use with caution!"""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("DELETE FROM memory_embeddings")
            conn.execute("DELETE FROM memory_chunks")
