"""
cortex/hippo/db/connection.py — aiosqlite connection management for Hippo.

Provides a ref-counted connection pool (single writer, concurrent readers)
with WAL mode and configurable busy_timeout.

Issue: #801
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Optional

import aiosqlite

from cortex.hippo.config import HippoConfig
from cortex.hippo.db.schema import get_full_schema_sql, HIPPO_SCHEMA_VERSION

logger = logging.getLogger(__name__)


class HippoDatabase:
    """Manages a single aiosqlite connection for Hippo.

    Opens the connection lazily on first use. The connection is shared
    across the engine lifetime (single writer + WAL = safe concurrent reads).
    """

    def __init__(self, config: HippoConfig) -> None:
        self._config = config
        self._conn: Optional[aiosqlite.Connection] = None
        self.vec_available: bool = False

    @property
    def conn(self) -> aiosqlite.Connection:
        """Return the active connection. Raises if not initialized."""
        if self._conn is None:
            raise RuntimeError("HippoDatabase not initialized — call initialize() first")
        return self._conn

    async def initialize(self) -> None:
        """Open the database connection, set pragmas, and create schema."""
        db_path = self._config.db_path
        Path(db_path).parent.mkdir(parents=True, exist_ok=True)

        self._conn = await aiosqlite.connect(db_path)
        self._conn.row_factory = aiosqlite.Row

        # Set pragmas
        if self._config.enable_wal:
            await self._conn.execute("PRAGMA journal_mode = WAL")
        await self._conn.execute("PRAGMA foreign_keys = ON")
        await self._conn.execute(
            f"PRAGMA busy_timeout = {self._config.busy_timeout_ms}"
        )
        await self._conn.execute(
            f"PRAGMA cache_size = {self._config.pragma_cache_size}"
        )
        await self._conn.execute(
            f"PRAGMA mmap_size = {self._config.pragma_mmap_size}"
        )

        # Load sqlite-vec extension if available and enabled
        if self._config.vec_enabled:
            await self._load_vec_extension()

        # Create schema (idempotent)
        await self._create_schema()

        logger.info(
            "Hippo database initialized: %s (WAL=%s, vec=%s)",
            db_path,
            self._config.enable_wal,
            self.vec_available,
        )

    async def _load_vec_extension(self) -> None:
        """Attempt to load the sqlite-vec extension. Graceful fallback on failure."""
        try:
            import sqlite_vec

            await self._conn.enable_load_extension(True)
            await self._conn.load_extension(sqlite_vec.loadable_path())
            await self._conn.enable_load_extension(False)
            self.vec_available = True
            logger.info("sqlite-vec extension loaded successfully")
        except ImportError:
            logger.warning(
                "sqlite-vec not installed — falling back to brute-force search. "
                "Install with: pip install sqlite-vec"
            )
        except Exception as exc:
            logger.warning("Failed to load sqlite-vec extension: %s", exc)

    async def _create_schema(self) -> None:
        """Create all Hippo tables, indexes, FTS, triggers, and vec0 if available."""
        from cortex.hippo.db.schema import ALL_DDL, get_vec_ddl

        for ddl_block in ALL_DDL:
            ddl_block = ddl_block.strip()
            if not ddl_block:
                continue
            try:
                await self._conn.executescript(ddl_block)
            except Exception as exc:
                # FTS/trigger creation may fail on some SQLite builds
                if "fts5" in ddl_block.lower() or "trigger" in ddl_block.lower():
                    logger.warning("FTS/trigger creation skipped: %s", exc)
                else:
                    raise

        # Create vec0 virtual table if extension is loaded
        if self.vec_available:
            try:
                vec_ddl = get_vec_ddl(self._config.embedding_dim)
                await self._conn.executescript(vec_ddl)
                logger.info("vec0 virtual table created (dim=%d)", self._config.embedding_dim)
            except Exception as exc:
                logger.warning("vec0 table creation failed: %s", exc)
                self.vec_available = False

        # Store schema version
        await self._conn.execute(
            """
            INSERT INTO hippo_config (key, value, updated_at)
            VALUES ('schema_version', ?, strftime('%Y-%m-%dT%H:%M:%SZ', 'now'))
            ON CONFLICT(key) DO UPDATE SET
                value = excluded.value,
                updated_at = excluded.updated_at
            """,
            (str(HIPPO_SCHEMA_VERSION),),
        )
        await self._conn.commit()

        # Run incremental migrations (backward-compat columns, etc.)
        await self._run_incremental_migrations()

        # Recover any stale compaction flags from previous crash
        from cortex.hippo.compaction.hardening import recover_stale_compaction_flags
        await recover_stale_compaction_flags(self._conn)

    async def _run_incremental_migrations(self) -> None:
        """Apply incremental migrations that add columns to existing tables."""
        from cortex.hippo.db.migrations.add_compaction_guard import migrate as _m_compaction_guard
        try:
            await _m_compaction_guard(self._conn)
        except Exception as exc:
            logger.warning("Incremental migration add_compaction_guard failed: %s", exc)

    async def close(self) -> None:
        """Close the database connection."""
        if self._conn:
            await self._conn.close()
            self._conn = None
            logger.info("Hippo database connection closed")

    async def get_schema_version(self) -> int:
        """Read the current schema version from hippo_config."""
        try:
            async with self.conn.execute(
                "SELECT value FROM hippo_config WHERE key = 'schema_version'"
            ) as cur:
                row = await cur.fetchone()
                return int(row[0]) if row else 0
        except Exception:
            return 0
