app/database.py

from __future__ import annotations

import re
import sqlite3
import uuid
from contextlib import contextmanager
from datetime import UTC, datetime
from pathlib import Path


MEMORY_PATTERNS = (
    (re.compile(r"\bmy name is ([a-zA-Z][a-zA-Z\s'-]{1,40})", re.I), "Name"),
    (re.compile(r"\bi am allergic to ([^.!,;\n]+)", re.I), "Allergy"),
    (re.compile(r"\bi have ([^.!,;\n]+)", re.I), "Condition"),
    (re.compile(r"\bi am taking ([^.!,;\n]+)", re.I), "Medication"),
    (re.compile(r"\bmy age is (\d{1,3})\b", re.I), "Age"),
)


class MemoryStore:
    def __init__(self, database_path: Path) -> None:
        self.database_path = Path(database_path)
        self.database_path.parent.mkdir(parents=True, exist_ok=True)
        self._initialize()

    @contextmanager
    def _connect(self):
        connection = sqlite3.connect(self.database_path)
        connection.row_factory = sqlite3.Row
        try:
            yield connection
            connection.commit()
        finally:
            connection.close()

    def _initialize(self) -> None:
        with self._connect() as connection:
            connection.executescript(
                """
                CREATE TABLE IF NOT EXISTS users (
                    user_id TEXT PRIMARY KEY,
                    created_at TEXT NOT NULL,
                    last_seen_at TEXT NOT NULL
                );

                CREATE TABLE IF NOT EXISTS messages (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id TEXT NOT NULL,
                    role TEXT NOT NULL CHECK(role IN ('user', 'assistant')),
                    content TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    FOREIGN KEY(user_id) REFERENCES users(user_id)
                );

                CREATE TABLE IF NOT EXISTS memory_notes (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id TEXT NOT NULL,
                    note TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    UNIQUE(user_id, note),
                    FOREIGN KEY(user_id) REFERENCES users(user_id)
                );
                """
            )

    def create_user(self) -> dict[str, str]:
        user_id = f"MB-{uuid.uuid4().hex[:12].upper()}"
        now = datetime.now(UTC).isoformat()
        with self._connect() as connection:
            connection.execute(
                "INSERT INTO users (user_id, created_at, last_seen_at) VALUES (?, ?, ?)",
                (user_id, now, now),
            )
        return {"user_id": user_id, "created_at": now}

    def ensure_user(self, user_id: str) -> dict[str, str] | None:
        with self._connect() as connection:
            row = connection.execute(
                "SELECT user_id, created_at FROM users WHERE user_id = ?",
                (user_id,),
            ).fetchone()
            if not row:
                return None

            connection.execute(
                "UPDATE users SET last_seen_at = ? WHERE user_id = ?",
                (datetime.now(UTC).isoformat(), user_id),
            )
            return {"user_id": row["user_id"], "created_at": row["created_at"]}

    def add_message(self, user_id: str, role: str, content: str) -> None:
        timestamp = datetime.now(UTC).isoformat()
        with self._connect() as connection:
            connection.execute(
                "INSERT INTO messages (user_id, role, content, created_at) VALUES (?, ?, ?, ?)",
                (user_id, role, content, timestamp),
            )
            connection.execute(
                "UPDATE users SET last_seen_at = ? WHERE user_id = ?",
                (timestamp, user_id),
            )

        if role == "user":
            self._extract_memory_notes(user_id, content)

    def _extract_memory_notes(self, user_id: str, content: str) -> None:
        extracted: set[str] = set()
        normalized = " ".join(content.split())
        for pattern, label in MEMORY_PATTERNS:
            for match in pattern.findall(normalized):
                value = match.strip(" .")
                if 2 <= len(value) <= 120:
                    extracted.add(f"{label}: {value}")

        if not extracted:
            return

        now = datetime.now(UTC).isoformat()
        with self._connect() as connection:
            connection.executemany(
                "INSERT OR IGNORE INTO memory_notes (user_id, note, created_at) VALUES (?, ?, ?)",
                [(user_id, note, now) for note in sorted(extracted)],
            )

    def get_history(self, user_id: str, limit: int | None = None) -> list[dict[str, str]]:
        query = """
            SELECT role, content, created_at
            FROM messages
            WHERE user_id = ?
            ORDER BY id DESC
        """
        params: list[str | int] = [user_id]
        if limit is not None:
            query += " LIMIT ?"
            params.append(limit)

        with self._connect() as connection:
            rows = connection.execute(query, params).fetchall()

        return [dict(row) for row in reversed(rows)]

    def get_memory_notes(self, user_id: str, limit: int = 10) -> list[str]:
        with self._connect() as connection:
            rows = connection.execute(
                """
                SELECT note
                FROM memory_notes
                WHERE user_id = ?
                ORDER BY id DESC
                LIMIT ?
                """,
                (user_id, limit),
            ).fetchall()
        return [row["note"] for row in rows]