"""
Episodic Memory — Event-based temporal memory store.

Records events with temporal context, importance scoring, and full-text search.
Backed by SQLite + FTS5 for zero-dependency, portable storage.

Schema:
    events(id, timestamp, source, author, content, importance,
           emotional_valence, tags, embedding_json)
"""

from __future__ import annotations

import json
import time
import uuid
from dataclasses import dataclass, field
from typing import Any

import aiosqlite


# ---------------------------------------------------------------------------
# Data models
# ---------------------------------------------------------------------------

@dataclass
class Episode:
    """A single episodic memory entry."""
    id: str
    timestamp: float
    source: str
    author: str
    content: str
    importance: float
    emotional_valence: float  # -1.0 (negative) to 1.0 (positive)
    tags: list[str]
    access_count: int = 0
    last_accessed: float = 0.0


@dataclass
class EpisodeSearchResult:
    """Episodic memory search result with relevance score."""
    episode: Episode
    relevance: float  # 0.0 – 1.0


# ---------------------------------------------------------------------------
# Episodic Memory Store
# ---------------------------------------------------------------------------

class EpisodicMemory:
    """
    SQLite + FTS5 backed episodic memory store.

    Key features:
    - Stores events with temporal context and importance scores
    - Full-text search via FTS5
    - Temporal range queries
    - Access tracking (strengthens frequently accessed memories)
    """

    def __init__(self, db_path: str = ":memory:") -> None:
        self._db_path = db_path
        self._db: aiosqlite.Connection | None = None

    async def initialize(self) -> None:
        """Initialize the database and create tables."""
        self._db = await aiosqlite.connect(self._db_path)
        await self._db.execute("PRAGMA journal_mode=WAL")
        await self._db.execute("PRAGMA foreign_keys=ON")

        await self._db.executescript("""
            CREATE TABLE IF NOT EXISTS episodes (
                id TEXT PRIMARY KEY,
                timestamp REAL NOT NULL,
                source TEXT NOT NULL DEFAULT 'unknown',
                author TEXT NOT NULL DEFAULT 'unknown',
                content TEXT NOT NULL,
                importance REAL NOT NULL DEFAULT 0.5,
                emotional_valence REAL NOT NULL DEFAULT 0.0,
                tags_json TEXT NOT NULL DEFAULT '[]',
                access_count INTEGER NOT NULL DEFAULT 0,
                last_accessed REAL NOT NULL DEFAULT 0.0,
                created_at REAL NOT NULL DEFAULT (unixepoch('now'))
            );

            CREATE INDEX IF NOT EXISTS idx_episodes_timestamp ON episodes(timestamp);
            CREATE INDEX IF NOT EXISTS idx_episodes_importance ON episodes(importance);

            CREATE VIRTUAL TABLE IF NOT EXISTS episodes_fts USING fts5(
                content,
                content=episodes,
                content_rowid=rowid
            );

            CREATE TRIGGER IF NOT EXISTS episodes_ai AFTER INSERT ON episodes BEGIN
                INSERT INTO episodes_fts(rowid, content)
                VALUES (new.rowid, new.content);
            END;

            CREATE TRIGGER IF NOT EXISTS episodes_ad AFTER DELETE ON episodes BEGIN
                INSERT INTO episodes_fts(episodes_fts, rowid, content)
                VALUES ('delete', old.rowid, old.content);
            END;

            CREATE TRIGGER IF NOT EXISTS episodes_au AFTER UPDATE ON episodes BEGIN
                INSERT INTO episodes_fts(episodes_fts, rowid, content)
                VALUES ('delete', old.rowid, old.content);
                INSERT INTO episodes_fts(rowid, content)
                VALUES (new.rowid, new.content);
            END;
        """)
        await self._db.commit()

    async def store(
        self,
        content: str,
        source: str = "conversation",
        author: str = "user",
        importance: float = 0.5,
        emotional_valence: float = 0.0,
        tags: list[str] | None = None,
    ) -> Episode:
        """Store a new episodic memory."""
        assert self._db is not None, "Call initialize() first"

        episode = Episode(
            id=uuid.uuid4().hex[:12],
            timestamp=time.time(),
            source=source,
            author=author,
            content=content,
            importance=importance,
            emotional_valence=emotional_valence,
            tags=tags or [],
        )

        await self._db.execute(
            """INSERT INTO episodes (id, timestamp, source, author, content,
               importance, emotional_valence, tags_json)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                episode.id,
                episode.timestamp,
                episode.source,
                episode.author,
                episode.content,
                episode.importance,
                episode.emotional_valence,
                json.dumps(episode.tags),
            ),
        )
        await self._db.commit()
        return episode

    async def search(
        self,
        query: str,
        limit: int = 10,
        min_importance: float = 0.0,
    ) -> list[EpisodeSearchResult]:
        """Full-text search across episodic memories."""
        assert self._db is not None

        rows = await self._db.execute_fetchall(
            """SELECT e.id, e.timestamp, e.source, e.author, e.content,
                      e.importance, e.emotional_valence, e.tags_json,
                      e.access_count, e.last_accessed,
                      rank
               FROM episodes_fts fts
               JOIN episodes e ON e.rowid = fts.rowid
               WHERE episodes_fts MATCH ?
               AND e.importance >= ?
               ORDER BY rank
               LIMIT ?""",
            (query, min_importance, limit),
        )

        results = []
        for row in rows:
            episode = self._row_to_episode(row)
            # FTS5 rank is negative (lower = more relevant), normalize to 0-1
            raw_rank = abs(row[10]) if row[10] else 1.0
            relevance = max(0.0, min(1.0, 1.0 / (1.0 + raw_rank)))
            results.append(EpisodeSearchResult(episode=episode, relevance=relevance))

        # Batch access tracking (single commit for all results)
        if results:
            for r in results:
                await self._track_access(r.episode.id)
            await self._db.commit()

        return results

    async def get_recent(
        self,
        limit: int = 20,
        since: float | None = None,
    ) -> list[Episode]:
        """Get recent episodes, optionally since a timestamp."""
        assert self._db is not None

        if since:
            rows = await self._db.execute_fetchall(
                """SELECT id, timestamp, source, author, content,
                          importance, emotional_valence, tags_json,
                          access_count, last_accessed
                   FROM episodes
                   WHERE timestamp >= ?
                   ORDER BY timestamp DESC
                   LIMIT ?""",
                (since, limit),
            )
        else:
            rows = await self._db.execute_fetchall(
                """SELECT id, timestamp, source, author, content,
                          importance, emotional_valence, tags_json,
                          access_count, last_accessed
                   FROM episodes
                   ORDER BY timestamp DESC
                   LIMIT ?""",
                (limit,),
            )

        return [self._row_to_episode(row) for row in rows]

    async def get_important(
        self,
        limit: int = 10,
        min_importance: float = 0.7,
    ) -> list[Episode]:
        """Get the most important memories."""
        assert self._db is not None

        rows = await self._db.execute_fetchall(
            """SELECT id, timestamp, source, author, content,
                      importance, emotional_valence, tags_json,
                      access_count, last_accessed
               FROM episodes
               WHERE importance >= ?
               ORDER BY importance DESC, access_count DESC
               LIMIT ?""",
            (min_importance, limit),
        )

        return [self._row_to_episode(row) for row in rows]

    async def count(self) -> int:
        """Total number of episodes stored."""
        assert self._db is not None
        row = await self._db.execute_fetchall("SELECT COUNT(*) FROM episodes")
        return row[0][0] if row else 0

    async def close(self) -> None:
        """Close the database connection."""
        if self._db:
            await self._db.close()
            self._db = None

    # -- Internal -----------------------------------------------------------

    async def _track_access(self, episode_id: str) -> None:
        """Update access count and last_accessed time."""
        if self._db:
            await self._db.execute(
                """UPDATE episodes
                   SET access_count = access_count + 1,
                       last_accessed = ?
                   WHERE id = ?""",
                (time.time(), episode_id),
            )

    def _row_to_episode(self, row: tuple) -> Episode:
        return Episode(
            id=row[0],
            timestamp=row[1],
            source=row[2],
            author=row[3],
            content=row[4],
            importance=row[5],
            emotional_valence=row[6],
            tags=json.loads(row[7]) if row[7] else [],
            access_count=row[8] if len(row) > 8 else 0,
            last_accessed=row[9] if len(row) > 9 else 0.0,
        )
