# core/store.py
from __future__ import annotations

import json
import sqlite3
from typing import Protocol


class StateStore(Protocol):
    async def get(self, key: str) -> dict | None: ...
    async def set(self, key: str, value: dict) -> None: ...
    async def delete(self, key: str) -> None: ...
    async def keys(self, prefix: str) -> list[str]: ...


class InMemoryStore:
    def __init__(self) -> None:
        self._data: dict[str, dict] = {}

    async def get(self, key: str) -> dict | None:
        return self._data.get(key)

    async def set(self, key: str, value: dict) -> None:
        self._data[key] = value

    async def delete(self, key: str) -> None:
        self._data.pop(key, None)

    async def keys(self, prefix: str) -> list[str]:
        return [k for k in self._data if k.startswith(prefix)]


class SQLiteStore:
    def __init__(self, db_path: str) -> None:
        self._db_path = db_path
        self._init_db()

    def _init_db(self) -> None:
        conn = sqlite3.connect(self._db_path)
        conn.execute(
            "CREATE TABLE IF NOT EXISTS kv (key TEXT PRIMARY KEY, value TEXT NOT NULL)"
        )
        conn.commit()
        conn.close()

    async def get(self, key: str) -> dict | None:
        conn = sqlite3.connect(self._db_path)
        row = conn.execute("SELECT value FROM kv WHERE key = ?", (key,)).fetchone()
        conn.close()
        return json.loads(row[0]) if row else None

    async def set(self, key: str, value: dict) -> None:
        conn = sqlite3.connect(self._db_path)
        conn.execute(
            "INSERT OR REPLACE INTO kv (key, value) VALUES (?, ?)",
            (key, json.dumps(value, default=str)),
        )
        conn.commit()
        conn.close()

    async def delete(self, key: str) -> None:
        conn = sqlite3.connect(self._db_path)
        conn.execute("DELETE FROM kv WHERE key = ?", (key,))
        conn.commit()
        conn.close()

    async def keys(self, prefix: str) -> list[str]:
        conn = sqlite3.connect(self._db_path)
        rows = conn.execute(
            "SELECT key FROM kv WHERE key LIKE ?", (prefix + "%",)
        ).fetchall()
        conn.close()
        return [row[0] for row in rows]