from __future__ import annotations

import json
import sqlite3
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator

from .schema import SCHEMA_SQL, SCHEMA_VERSION


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

    def initialize(self) -> None:
        with self.connection() as conn:
            conn.executescript(SCHEMA_SQL)
            self._ensure_column(conn, "memories", "schema_version", f"TEXT NOT NULL DEFAULT '{SCHEMA_VERSION}'")
            self._ensure_column(conn, "memories", "subtype", "TEXT")
            self._ensure_column(conn, "memories", "run_id", "TEXT")
            self._ensure_column(conn, "memories", "task_id", "TEXT")
            self._ensure_column(conn, "memories", "url", "TEXT")
            self._ensure_column(conn, "memories", "domain", "TEXT")
            self._ensure_column(conn, "memories", "origin_agent", "TEXT")
            self._ensure_column(conn, "tasks", "schema_version", f"TEXT NOT NULL DEFAULT '{SCHEMA_VERSION}'")
            self._ensure_column(conn, "tasks", "run_id", "TEXT")

    @contextmanager
    def connection(self) -> Iterator[sqlite3.Connection]:
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA foreign_keys = ON")
        try:
            yield conn
            conn.commit()
        finally:
            conn.close()

    @staticmethod
    def dumps(value: object) -> str:
        return json.dumps(value, ensure_ascii=True, sort_keys=True)

    @staticmethod
    def loads(value: str | None) -> object:
        if not value:
            return None
        return json.loads(value)

    @staticmethod
    def _ensure_column(conn: sqlite3.Connection, table: str, column: str, definition: str) -> None:
        existing = {row["name"] for row in conn.execute(f"PRAGMA table_info({table})")}
        if column in existing:
            return
        conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {definition}")