"""Storage — SQLite persistence layer for XYGLUE data.

Provides async SQLite storage for patterns, templates, ground records,
and run metadata. Uses aiosqlite for non-blocking I/O.

Supports:
- JSON document storage (save/load/list/search/delete)
- Table-specific queries with filtering
- Run history tracking
- Schema migrations
- Health checks and stats
- Context manager protocol for clean lifecycle
"""

import json
import time
from typing import Any

import aiosqlite


# ── Table names (used as constants to avoid SQL injection) ─────────

VALID_TABLES = {"patterns", "templates", "ground_records", "runs"}


class Storage:
    """SQLite-based persistence for patterns, templates, and ground records.

    Usage:
        storage = Storage("xyglue.db")
        await storage.initialize()
        # ... use storage ...
        await storage.close()

    Or with async context manager:
        async with Storage("xyglue.db") as storage:
            await storage.save_json("patterns", "id1", {...})
    """

    SCHEMA_VERSION = 1

    def __init__(self, db_path: str = "xyglue.db") -> None:
        self.db_path = db_path
        self._db: aiosqlite.Connection | None = None
        self._initialized: bool = False

    # ── Lifecycle ──────────────────────────────────────────────────

    async def initialize(self) -> None:
        """Open database connection and create tables if needed."""
        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._create_tables()
        self._initialized = True

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

    async def __aenter__(self) -> "Storage":
        await self.initialize()
        return self

    async def __aexit__(self, *args: Any) -> None:
        await self.close()

    @property
    def is_initialized(self) -> bool:
        return self._initialized and self._db is not None

    # ── Schema ─────────────────────────────────────────────────────

    async def _create_tables(self) -> None:
        """Create all tables if they don't exist."""
        if not self._db:
            return
        await self._db.executescript("""
            CREATE TABLE IF NOT EXISTS patterns (
                id TEXT PRIMARY KEY,
                name TEXT,
                category TEXT,
                data TEXT,
                created_at REAL
            );
            CREATE INDEX IF NOT EXISTS idx_patterns_category ON patterns(category);
            CREATE INDEX IF NOT EXISTS idx_patterns_name ON patterns(name);

            CREATE TABLE IF NOT EXISTS templates (
                id TEXT PRIMARY KEY,
                domain TEXT,
                quality REAL DEFAULT 0.0,
                data TEXT,
                created_at REAL
            );
            CREATE INDEX IF NOT EXISTS idx_templates_domain ON templates(domain);
            CREATE INDEX IF NOT EXISTS idx_templates_quality ON templates(quality);

            CREATE TABLE IF NOT EXISTS ground_records (
                id TEXT PRIMARY KEY,
                run_id TEXT,
                domain TEXT,
                success INTEGER DEFAULT 0,
                data TEXT,
                created_at REAL
            );
            CREATE INDEX IF NOT EXISTS idx_ground_records_run ON ground_records(run_id);
            CREATE INDEX IF NOT EXISTS idx_ground_records_domain ON ground_records(domain);

            CREATE TABLE IF NOT EXISTS runs (
                id TEXT PRIMARY KEY,
                goal TEXT,
                status TEXT,
                mode TEXT,
                data TEXT,
                created_at REAL,
                completed_at REAL
            );
            CREATE INDEX IF NOT EXISTS idx_runs_status ON runs(status);

            CREATE TABLE IF NOT EXISTS schema_version (
                version INTEGER PRIMARY KEY
            );
        """)
        # Track schema version
        async with self._db.execute(
            "SELECT version FROM schema_version ORDER BY version DESC LIMIT 1"
        ) as cursor:
            row = await cursor.fetchone()
            if not row:
                await self._db.execute(
                    "INSERT INTO schema_version (version) VALUES (?)",
                    [self.SCHEMA_VERSION],
                )
        await self._db.commit()

    # ── Generic JSON CRUD ──────────────────────────────────────────

    def _validate_table(self, table: str) -> None:
        """Validate table name to prevent SQL injection."""
        if table not in VALID_TABLES:
            raise ValueError(f"Invalid table name: {table}")

    async def save_json(
        self, table: str, id: str, data: dict[str, Any], **extra: Any
    ) -> None:
        """Save a JSON document to a table."""
        self._validate_table(table)
        if not self._db:
            return
        cols = ["id", "data"] + list(extra.keys())
        vals = [id, json.dumps(data, default=str)] + list(extra.values())
        placeholders = ",".join(["?"] * len(vals))
        col_str = ",".join(cols)
        await self._db.execute(
            f"INSERT OR REPLACE INTO {table} ({col_str}) VALUES ({placeholders})",
            vals,
        )
        await self._db.commit()

    async def load_json(self, table: str, id: str) -> dict[str, Any] | None:
        """Load a JSON document by ID."""
        self._validate_table(table)
        if not self._db:
            return None
        async with self._db.execute(
            f"SELECT data FROM {table} WHERE id = ?", [id]
        ) as cursor:
            row = await cursor.fetchone()
            return json.loads(row[0]) if row else None

    async def delete(self, table: str, id: str) -> bool:
        """Delete a document by ID. Returns True if a row was deleted."""
        self._validate_table(table)
        if not self._db:
            return False
        cursor = await self._db.execute(
            f"DELETE FROM {table} WHERE id = ?", [id]
        )
        await self._db.commit()
        return cursor.rowcount > 0

    async def exists(self, table: str, id: str) -> bool:
        """Check if a document exists."""
        self._validate_table(table)
        if not self._db:
            return False
        async with self._db.execute(
            f"SELECT 1 FROM {table} WHERE id = ? LIMIT 1", [id]
        ) as cursor:
            row = await cursor.fetchone()
            return row is not None

    # ── List and search ────────────────────────────────────────────

    async def list_ids(self, table: str, limit: int = 100) -> list[str]:
        """List document IDs in a table."""
        self._validate_table(table)
        if not self._db:
            return []
        async with self._db.execute(
            f"SELECT id FROM {table} ORDER BY created_at DESC LIMIT ?", [limit]
        ) as cursor:
            rows = await cursor.fetchall()
            return [row[0] for row in rows]

    async def list_all(
        self, table: str, limit: int = 100
    ) -> list[dict[str, Any]]:
        """List all documents in a table as parsed JSON dicts."""
        self._validate_table(table)
        if not self._db:
            return []
        async with self._db.execute(
            f"SELECT id, data FROM {table} ORDER BY created_at DESC LIMIT ?",
            [limit],
        ) as cursor:
            rows = await cursor.fetchall()
            results: list[dict[str, Any]] = []
            for row in rows:
                try:
                    doc = json.loads(row[1])
                    doc["_id"] = row[0]
                    results.append(doc)
                except json.JSONDecodeError:
                    continue
            return results

    async def count(self, table: str) -> int:
        """Count documents in a table."""
        self._validate_table(table)
        if not self._db:
            return 0
        async with self._db.execute(f"SELECT COUNT(*) FROM {table}") as cursor:
            row = await cursor.fetchone()
            return row[0] if row else 0

    # ── Filtered queries ───────────────────────────────────────────

    async def find_by_field(
        self, table: str, field: str, value: Any, limit: int = 50
    ) -> list[dict[str, Any]]:
        """Find documents where a column equals a value.

        Only searches indexed columns (not inside JSON data).
        """
        self._validate_table(table)
        if not self._db:
            return []
        # Only allow querying known columns to prevent injection
        allowed_columns = {
            "patterns": {"name", "category"},
            "templates": {"domain"},
            "ground_records": {"run_id", "domain"},
            "runs": {"status", "mode"},
        }
        if field not in allowed_columns.get(table, set()):
            raise ValueError(f"Cannot query field '{field}' on table '{table}'")
        async with self._db.execute(
            f"SELECT id, data FROM {table} WHERE {field} = ? LIMIT ?",
            [value, limit],
        ) as cursor:
            rows = await cursor.fetchall()
            results: list[dict[str, Any]] = []
            for row in rows:
                try:
                    doc = json.loads(row[1])
                    doc["_id"] = row[0]
                    results.append(doc)
                except json.JSONDecodeError:
                    continue
            return results

    # ── Run tracking ───────────────────────────────────────────────

    async def save_run(
        self,
        run_id: str,
        goal: str,
        status: str,
        mode: str,
        data: dict[str, Any],
    ) -> None:
        """Save or update a run record."""
        if not self._db:
            return
        now = time.time()
        completed = now if status in ("completed", "failed") else None
        await self._db.execute(
            """INSERT OR REPLACE INTO runs
               (id, goal, status, mode, data, created_at, completed_at)
               VALUES (?, ?, ?, ?, ?, ?, ?)""",
            [run_id, goal, status, mode, json.dumps(data, default=str), now, completed],
        )
        await self._db.commit()

    async def get_run(self, run_id: str) -> dict[str, Any] | None:
        """Get a run record by ID."""
        if not self._db:
            return None
        async with self._db.execute(
            "SELECT id, goal, status, mode, data, created_at, completed_at FROM runs WHERE id = ?",
            [run_id],
        ) as cursor:
            row = await cursor.fetchone()
            if not row:
                return None
            data = json.loads(row[4]) if row[4] else {}
            return {
                "id": row[0],
                "goal": row[1],
                "status": row[2],
                "mode": row[3],
                "data": data,
                "created_at": row[5],
                "completed_at": row[6],
            }

    async def list_runs(
        self, status: str | None = None, limit: int = 50
    ) -> list[dict[str, Any]]:
        """List runs, optionally filtered by status."""
        if not self._db:
            return []
        if status:
            query = "SELECT id, goal, status, mode, created_at FROM runs WHERE status = ? ORDER BY created_at DESC LIMIT ?"
            params = [status, limit]
        else:
            query = "SELECT id, goal, status, mode, created_at FROM runs ORDER BY created_at DESC LIMIT ?"
            params = [limit]
        async with self._db.execute(query, params) as cursor:
            rows = await cursor.fetchall()
            return [
                {
                    "id": row[0],
                    "goal": row[1],
                    "status": row[2],
                    "mode": row[3],
                    "created_at": row[4],
                }
                for row in rows
            ]

    # ── Health and stats ───────────────────────────────────────────

    async def health_check(self) -> dict[str, Any]:
        """Run a health check on the storage layer."""
        if not self._db:
            return {"healthy": False, "error": "Not initialized"}
        try:
            async with self._db.execute("SELECT 1") as cursor:
                await cursor.fetchone()
            counts: dict[str, int] = {}
            for table in VALID_TABLES:
                counts[table] = await self.count(table)
            return {
                "healthy": True,
                "db_path": self.db_path,
                "tables": counts,
            }
        except Exception as e:
            return {"healthy": False, "error": str(e)}

    async def get_stats(self) -> dict[str, Any]:
        """Get storage statistics."""
        if not self._db:
            return {"initialized": False}
        counts: dict[str, int] = {}
        for table in VALID_TABLES:
            counts[table] = await self.count(table)
        return {
            "initialized": True,
            "db_path": self.db_path,
            "tables": counts,
            "total_documents": sum(counts.values()),
        }

    # ── Cleanup ────────────────────────────────────────────────────

    async def cleanup_old(
        self, table: str, max_age_days: float = 90
    ) -> int:
        """Delete documents older than max_age_days. Returns count deleted."""
        self._validate_table(table)
        if not self._db:
            return 0
        cutoff = time.time() - (max_age_days * 86400)
        cursor = await self._db.execute(
            f"DELETE FROM {table} WHERE created_at < ?", [cutoff]
        )
        await self._db.commit()
        return cursor.rowcount
