"""SQLite-backed episode memory."""

from __future__ import annotations

import json
import re
import sqlite3
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

from .trace import Episode


_TOKEN_RE = re.compile(r"[A-Za-z0-9_]+")


def _tokens(text: str) -> set[str]:
    return {item.lower() for item in _TOKEN_RE.findall(text)}


@dataclass(slots=True)
class MemoryHit:
    episode: Episode
    score: float


class EpisodeMemory:
    """Local episode store.

    Retrieval is intentionally simple for the MVP. Later versions can add vector
    indexes, BM25, or hybrid retrieval without changing the Episode schema.
    """

    def __init__(self, path: str | Path) -> None:
        self.path = Path(path)
        if self.path.parent:
            self.path.parent.mkdir(parents=True, exist_ok=True)
        self._init_db()

    def _connect(self) -> sqlite3.Connection:
        con = sqlite3.connect(self.path)
        con.row_factory = sqlite3.Row
        return con

    def _init_db(self) -> None:
        with self._connect() as con:
            con.execute(
                """
                CREATE TABLE IF NOT EXISTS episodes (
                    episode_id TEXT PRIMARY KEY,
                    task TEXT NOT NULL,
                    outcome TEXT NOT NULL,
                    score REAL NOT NULL,
                    tags_json TEXT NOT NULL,
                    episode_json TEXT NOT NULL,
                    created_at TEXT NOT NULL
                )
                """
            )
            con.execute("CREATE INDEX IF NOT EXISTS idx_episodes_task ON episodes(task)")
            con.execute("CREATE INDEX IF NOT EXISTS idx_episodes_outcome ON episodes(outcome)")

    def add_episode(self, episode: Episode) -> None:
        payload = json.dumps(episode.to_dict(), sort_keys=True)
        tags_json = json.dumps(episode.tags, sort_keys=True)
        with self._connect() as con:
            con.execute(
                """
                INSERT OR REPLACE INTO episodes
                (episode_id, task, outcome, score, tags_json, episode_json, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    episode.episode_id,
                    episode.task,
                    episode.outcome,
                    episode.score,
                    tags_json,
                    payload,
                    episode.created_at,
                ),
            )

    def get_episode(self, episode_id: str) -> Episode | None:
        with self._connect() as con:
            row = con.execute("SELECT episode_json FROM episodes WHERE episode_id = ?", (episode_id,)).fetchone()
        if row is None:
            return None
        return Episode.from_dict(json.loads(str(row["episode_json"])))

    def all_episodes(self) -> list[Episode]:
        with self._connect() as con:
            rows = con.execute("SELECT episode_json FROM episodes ORDER BY created_at ASC").fetchall()
        return [Episode.from_dict(json.loads(str(row["episode_json"]))) for row in rows]

    def search(self, query: str, limit: int = 5, outcomes: Iterable[str] | None = None) -> list[MemoryHit]:
        query_tokens = _tokens(query)
        allowed_outcomes = set(outcomes) if outcomes is not None else None
        hits: list[MemoryHit] = []
        for episode in self.all_episodes():
            if allowed_outcomes is not None and episode.outcome not in allowed_outcomes:
                continue
            haystack = " ".join(
                [
                    episode.task,
                    episode.outcome,
                    " ".join(episode.tags),
                    json.dumps(episode.metadata, sort_keys=True),
                    " ".join(step.result_summary for step in episode.steps),
                ]
            )
            hay_tokens = _tokens(haystack)
            overlap = len(query_tokens & hay_tokens)
            if overlap:
                score = overlap / max(1, len(query_tokens))
                if episode.outcome == "success":
                    score += 0.05
                hits.append(MemoryHit(episode=episode, score=score))
        hits.sort(key=lambda item: item.score, reverse=True)
        return hits[:limit]

    def statistics(self) -> dict[str, int]:
        with self._connect() as con:
            rows = con.execute("SELECT outcome, COUNT(*) AS count FROM episodes GROUP BY outcome").fetchall()
        stats = {str(row["outcome"]): int(row["count"]) for row in rows}
        stats["total"] = sum(stats.values())
        return stats
