"""Database layer for GreyMatter.

Provides a Protocol-based database abstraction with SQLite and CockroachDB
implementations. The SQLite schema mirrors CockroachDB table structure for
clean upgrade path. Solo mode uses SQLite; Distributed mode uses CockroachDB.

AsyncDatabase wraps any Database with asyncio-compatible methods via
run_in_executor, allowing MCP tools and other async code to call the
blocking database without stalling the event loop.
"""
import asyncio
import logging
import os
import re
import sqlite3
import threading
from contextlib import asynccontextmanager, contextmanager
from pathlib import Path
from typing import Any, Callable, Protocol, runtime_checkable

import sqlite_vec

# Optional psycopg2 for CockroachDB support
try:
    import psycopg2
    from psycopg2 import pool as pg_pool
    from psycopg2.extras import RealDictCursor
    PSYCOPG2_AVAILABLE = True
except ImportError:
    PSYCOPG2_AVAILABLE = False

logger = logging.getLogger(__name__)


def derive_key_from_password(password: str, salt: bytes | None = None) -> tuple[str, bytes]:
    """Derive a hex encryption key from a passphrase using PBKDF2-HMAC-SHA256.

    Uses the FIPS-aware wrapper (:func:`~greymatter_core.security.fips.fips_pbkdf2`
    and :func:`~greymatter_core.security.fips.fips_random_bytes`) so that key
    derivation routes through the OpenSSL FIPS provider when ``GM_FIPS_MODE=1``.

    Args:
        password: Human-readable passphrase.
        salt: 16-byte salt. Generated via ``fips_random_bytes(16)`` if *None*.

    Returns:
        Tuple of (hex_key, salt) — store the salt alongside the DB for
        re-derivation.  The key is 256 bits (64 hex characters), suitable
        for AES-256 / SQLCipher.
    """
    from greymatter_core.security.fips import fips_pbkdf2, fips_random_bytes
    if salt is None:
        salt = fips_random_bytes(16)
    # 256,000 iterations matches our SQLCipher PRAGMA kdf_iter setting
    dk = fips_pbkdf2(password.encode("utf-8"), salt, iterations=256_000)
    return dk.hex(), salt


def _sqlcipher_row_factory(cursor, row):
    """Row factory for pysqlcipher3 that returns dict-like objects.

    pysqlcipher3 doesn't support sqlite3.Row, so we create our own
    dict wrapper that supports both dict-style and index-style access.
    """
    columns = [desc[0] for desc in cursor.description]
    return dict(zip(columns, row))

# Schema extension registry — extensions register callbacks before DB init
_schema_extensions: list[Callable[[sqlite3.Connection], None]] = []


def register_schema_extension(fn: Callable[[sqlite3.Connection], None]) -> None:
    """Register a schema extension callback.

    Called during SQLiteDatabase.__init__ before migrate(). The callback
    receives the raw sqlite3 connection and should create tables/triggers/indexes.
    Must be called BEFORE SQLiteDatabase is instantiated.
    """
    _schema_extensions.append(fn)


CORE_SCHEMA = """
CREATE TABLE IF NOT EXISTS patterns (
    id TEXT PRIMARY KEY,
    type TEXT NOT NULL,
    domain TEXT NOT NULL,
    title TEXT NOT NULL,
    observation TEXT NOT NULL,
    confidence TEXT DEFAULT 'low',
    confirmations INTEGER DEFAULT 0,
    applications INTEGER DEFAULT 0,
    last_applied_at TEXT,
    created_at TEXT DEFAULT (datetime('now')),
    updated_at TEXT DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS pattern_evidence (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    pattern_id TEXT NOT NULL REFERENCES patterns(id),
    session_id TEXT,
    work_item_id TEXT,
    outcome TEXT DEFAULT 'confirmed',
    created_at TEXT DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS knowledge_entries (
    id TEXT PRIMARY KEY,
    entry_type TEXT NOT NULL,
    scope TEXT DEFAULT 'cross-cutting',
    domain TEXT,
    title TEXT NOT NULL,
    content TEXT NOT NULL,
    valence TEXT,
    confidence REAL DEFAULT 0.5,
    applied_count INTEGER DEFAULT 0,
    tags TEXT,
    classification_level INTEGER NOT NULL DEFAULT 3,
    created_at TEXT DEFAULT (datetime('now')),
    updated_at TEXT DEFAULT (datetime('now')),
    pruned_at TEXT,
    stale_flagged_at TEXT,
    last_confirmed_at TEXT,
    superseded_by TEXT
);

CREATE TABLE IF NOT EXISTS work_items (
    id TEXT PRIMARY KEY,
    title TEXT NOT NULL,
    description TEXT,
    status TEXT DEFAULT 'queued',
    priority INTEGER DEFAULT 100,
    domain_tags TEXT,
    assigned_to TEXT,
    project TEXT DEFAULT 'greymatter',
    feature TEXT DEFAULT 'general',
    type TEXT DEFAULT 'task',
    short_id TEXT,
    created_at TEXT DEFAULT (datetime('now')),
    completed_at TEXT,
    result_summary TEXT
);

CREATE TABLE IF NOT EXISTS sessions (
    id TEXT PRIMARY KEY,
    started_at TEXT DEFAULT (datetime('now')),
    ended_at TEXT,
    project TEXT,
    patterns_applied TEXT,
    work_completed TEXT,
    summary TEXT,
    state TEXT DEFAULT 'active',
    active_work_item_id TEXT,
    checkpoint_at TEXT,
    compaction_count INTEGER DEFAULT 0,
    environment TEXT
);

CREATE TABLE IF NOT EXISTS session_checkpoints (
    id TEXT PRIMARY KEY,
    session_id TEXT NOT NULL,
    checkpoint_at TEXT NOT NULL,
    state TEXT NOT NULL,
    active_work_item_id TEXT,
    compaction_count INTEGER DEFAULT 0,
    context_summary TEXT,
    open_items TEXT,
    FOREIGN KEY (session_id) REFERENCES sessions(id)
);

CREATE TABLE IF NOT EXISTS llm_completions (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    timestamp TEXT DEFAULT (datetime('now')),
    provider TEXT NOT NULL,
    model TEXT NOT NULL,
    tokens_used INTEGER,
    cost_estimate REAL DEFAULT 0.0,
    task_type TEXT,
    latency_ms INTEGER
);

CREATE TABLE IF NOT EXISTS quality_standards (
    id TEXT PRIMARY KEY,
    category TEXT NOT NULL,
    rule TEXT NOT NULL,
    severity TEXT NOT NULL,
    example_good TEXT,
    example_bad TEXT,
    rationale TEXT,
    active INTEGER DEFAULT 1,
    created_at TEXT DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS build_runs (
    id TEXT PRIMARY KEY,
    session_id TEXT REFERENCES sessions(id),
    approach TEXT NOT NULL,
    task_name TEXT NOT NULL,
    started_at TEXT DEFAULT (datetime('now')),
    completed_at TEXT,
    lines_written INTEGER,
    test_count INTEGER,
    tests_passing INTEGER,
    first_pass_success INTEGER DEFAULT 0
);

CREATE TABLE IF NOT EXISTS quality_scores (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    build_run_id TEXT NOT NULL REFERENCES build_runs(id),
    standard_id TEXT NOT NULL REFERENCES quality_standards(id),
    passed INTEGER NOT NULL,
    detail TEXT,
    found_by TEXT,
    found_at TEXT DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS review_issues (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    build_run_id TEXT NOT NULL REFERENCES build_runs(id),
    severity TEXT NOT NULL,
    category TEXT NOT NULL,
    description TEXT NOT NULL,
    fixed_at TEXT,
    time_to_fix_minutes INTEGER,
    found_by TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS documents (
    id TEXT PRIMARY KEY,
    doc_type TEXT NOT NULL,
    title TEXT NOT NULL,
    content TEXT NOT NULL,
    format TEXT DEFAULT 'markdown',
    tags TEXT,
    project TEXT,
    parent_id TEXT,
    summary TEXT,
    token_estimate INTEGER,
    created_at TEXT DEFAULT (datetime('now')),
    updated_at TEXT DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS context_requests (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    session_id TEXT,
    query TEXT NOT NULL,
    documents_served TEXT,
    tokens_served INTEGER,
    relevance_score REAL,
    requested_at TEXT DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS embedding_chunks (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    source_table TEXT NOT NULL,
    source_id TEXT NOT NULL,
    chunk_index INTEGER NOT NULL,
    chunk_text TEXT NOT NULL,
    created_at TEXT DEFAULT (datetime('now'))
);
"""

# Backward compatibility alias
SCHEMA = CORE_SCHEMA


@runtime_checkable
class Database(Protocol):
    """Abstract database interface.

    Implementations must support parameterized queries returning dicts.
    SQLite uses ? placeholders. CockroachDB uses %s placeholders.
    """

    @property
    def dialect(self) -> str:
        """Return 'sqlite' or 'crdb' to enable SQL dialect switching."""
        ...

    def execute(self, query: str, params: tuple = ()) -> list[dict[str, Any]]: ...
    def execute_one(self, query: str, params: tuple = ()) -> dict[str, Any] | None: ...
    def transaction(self): ...
    def migrate(self) -> None: ...


def now_sql(db: Database) -> str:
    """Return the SQL expression for current timestamp for the given backend."""
    if hasattr(db, "dialect") and db.dialect == "crdb":
        return "now()"
    return "datetime('now')"


_shared_db_instance: "SQLiteDatabase | None" = None
_shared_db_lock = __import__("threading").Lock()


def get_shared_db(db_path: str | None = None) -> "SQLiteDatabase":
    """Get or create a singleton SQLiteDatabase instance.

    Thread-safe. First call with a path sets the DB location.
    Subsequent calls return the same instance (ignoring path if different).
    """
    global _shared_db_instance
    with _shared_db_lock:
        if _shared_db_instance is None:
            _shared_db_instance = SQLiteDatabase(path=db_path or "~/.greymatter/greymatter.db")
        return _shared_db_instance


def _reset_shared_db() -> None:
    """Reset the singleton for testing. Not for production use."""
    global _shared_db_instance
    with _shared_db_lock:
        _shared_db_instance = None


class SQLiteDatabase:
    """SQLite database for solo mode.

    Single file at ~/.greymatter/greymatter.db (or :memory: for tests).

    When encryption_key is provided (hex string), uses SQLCipher for
    AES-256 encryption at rest. Requires pysqlcipher3 to be installed.
    """

    @property
    def dialect(self) -> str:
        return "sqlite"

    def __init__(
        self,
        path: str = "~/.greymatter/greymatter.db",
        encryption_key: str | None = None,
        encryption_password: str | None = None,
    ):
        if path == ":memory:":
            self._path = path
        else:
            self._path = str(Path(path).expanduser())
            parent = Path(self._path).parent
            parent.mkdir(parents=True, exist_ok=True)
            # Restrict directory to owner-only (rwx------)
            try:
                parent.chmod(0o700)
            except OSError:
                pass

        self._encrypted = False

        # Derive hex key from password if provided (password takes precedence)
        if encryption_password and not encryption_key:
            salt_path = Path(self._path + ".salt") if self._path != ":memory:" else None
            existing_salt = None
            if salt_path and salt_path.exists():
                existing_salt = salt_path.read_bytes()
            encryption_key, salt = derive_key_from_password(
                encryption_password, existing_salt
            )
            # Persist salt next to DB for future derivation
            if salt_path and not salt_path.exists():
                salt_path.write_bytes(salt)
                try:
                    salt_path.chmod(0o600)
                except OSError:
                    pass
            logger.info("Derived encryption key from password via PBKDF2")

        if encryption_key:
            # Validate encryption key is strictly hex before interpolation.
            # PRAGMA key does not support parameterized queries, so we must
            # validate the input to prevent injection via crafted key values.
            import re as _re
            if not _re.fullmatch(r"[0-9a-fA-F]+", encryption_key):
                raise ValueError(
                    "Encryption key must be a hex string (0-9, a-f, A-F only)"
                )

            # Try sqlcipher3 first (Python 3.14+ compatible), then pysqlcipher3
            sqlcipher_mod = None
            try:
                import sqlcipher3 as sqlcipher_mod
            except ImportError:
                try:
                    from pysqlcipher3 import dbapi2 as sqlcipher_mod
                except ImportError:
                    pass

            if sqlcipher_mod is not None:
                try:
                    self._conn = sqlcipher_mod.connect(
                        self._path, check_same_thread=False, timeout=30
                    )
                    self._conn.execute(f"PRAGMA key = \"x'{encryption_key}'\";")
                    self._conn.execute("PRAGMA cipher_page_size = 4096;")
                    self._conn.execute("PRAGMA kdf_iter = 256000;")
                    self._conn.execute("PRAGMA cipher_hmac_algorithm = HMAC_SHA256;")
                    # Verify the key works (fails on wrong key or unencrypted DB)
                    self._conn.execute("SELECT count(*) FROM sqlite_master;")
                    self._conn.row_factory = _sqlcipher_row_factory
                    self._encrypted = True
                    logger.info("SQLCipher encryption at rest: enabled")
                except Exception as e:
                    # Existing unencrypted DB — can't open with key
                    self._conn.close()
                    logger.warning(
                        "Database exists but is not encrypted — opening without "
                        "encryption. Run 'greymatter migrate-encrypt' to encrypt "
                        "an existing database. Error: %s", e
                    )
                    self._conn = sqlite3.connect(
                        self._path, check_same_thread=False, timeout=30
                    )
                    self._conn.row_factory = sqlite3.Row
            else:
                logger.warning(
                    "Neither sqlcipher3 nor pysqlcipher3 installed — "
                    "database will NOT be encrypted at rest. "
                    "Install with: pip install sqlcipher3"
                )
                self._conn = sqlite3.connect(
                    self._path, check_same_thread=False, timeout=30
                )
                self._conn.row_factory = sqlite3.Row
        else:
            self._conn = sqlite3.connect(
                self._path, check_same_thread=False, timeout=30
            )
            self._conn.row_factory = sqlite3.Row

        self._conn.execute("PRAGMA journal_mode=WAL")
        self._conn.execute("PRAGMA busy_timeout=30000")
        self._conn.execute("PRAGMA foreign_keys=ON")

        if not self._encrypted:
            # sqlite_vec extension loading only works with standard sqlite3
            self._conn.enable_load_extension(True)
            sqlite_vec.load(self._conn)
            self._conn.enable_load_extension(False)

        # Restrict DB file to owner-only (rw-------)
        if self._path != ":memory:":
            try:
                Path(self._path).chmod(0o600)
                # Also lock down WAL and SHM files if they exist
                for suffix in ["-wal", "-shm"]:
                    wal_path = Path(self._path + suffix)
                    if wal_path.exists():
                        wal_path.chmod(0o600)
            except OSError:
                pass

        self._in_transaction = False
        # Run registered schema extensions before migrate
        for ext_fn in _schema_extensions:
            ext_fn(self._conn)
        self.migrate()

    def migrate(self) -> None:
        """Run schema migrations."""
        self._conn.executescript(CORE_SCHEMA)
        # Session state columns (for harness continuity)
        for col, col_type in [
            ("state", "TEXT DEFAULT 'active'"),
            ("active_work_item_id", "TEXT"),
            ("checkpoint_at", "TEXT"),
            ("compaction_count", "INTEGER DEFAULT 0"),
            ("environment", "TEXT"),
        ]:
            try:
                self._conn.execute(
                    f"ALTER TABLE sessions ADD COLUMN {col} {col_type}"
                )
            except Exception:
                pass  # Column already exists (catches both sqlite3 and sqlcipher3 errors)

        # Observational memory
        self._conn.executescript("""
            CREATE TABLE IF NOT EXISTS observations (
                id TEXT PRIMARY KEY,
                session_id TEXT NOT NULL,
                trigger TEXT NOT NULL,
                category TEXT NOT NULL,
                title TEXT NOT NULL,
                content TEXT NOT NULL,
                raw_context TEXT,
                confidence REAL DEFAULT 0.5,
                stability_days REAL DEFAULT 45,
                last_confirmed_at TEXT,
                token_count INTEGER DEFAULT 0,
                context_usage_pct REAL DEFAULT 0.0,
                knowledge_id TEXT,
                pattern_id TEXT,
                source_tool TEXT DEFAULT '',
                domains TEXT DEFAULT '',
                created_at TEXT DEFAULT (datetime('now')),
                updated_at TEXT DEFAULT (datetime('now')),
                pruned INTEGER DEFAULT 0
            );
            CREATE INDEX IF NOT EXISTS idx_observations_session ON observations(session_id);
            CREATE INDEX IF NOT EXISTS idx_observations_category ON observations(category);
            CREATE INDEX IF NOT EXISTS idx_observations_confidence ON observations(confidence);
            CREATE INDEX IF NOT EXISTS idx_observations_pruned ON observations(pruned);

            CREATE VIRTUAL TABLE IF NOT EXISTS observations_fts USING fts5(
                title, content
            );

            CREATE TABLE IF NOT EXISTS observation_sessions (
                id TEXT PRIMARY KEY,
                session_id TEXT NOT NULL,
                trigger TEXT NOT NULL,
                context_usage_pct REAL DEFAULT 0.0,
                observations_extracted INTEGER DEFAULT 0,
                observations_deduplicated INTEGER DEFAULT 0,
                observations_stored INTEGER DEFAULT 0,
                extraction_method TEXT DEFAULT '',
                extraction_time_ms INTEGER DEFAULT 0,
                created_at TEXT DEFAULT (datetime('now'))
            );
        """)

        # Ensure observations has stability_days + last_confirmed_at (memory decay model)
        for col, col_type in [
            ("stability_days", "REAL DEFAULT 45"),
            ("last_confirmed_at", "TEXT"),
        ]:
            try:
                self._conn.execute(
                    f"ALTER TABLE observations ADD COLUMN {col} {col_type}"
                )
            except Exception:  # Column/table may already exist (sqlite3 or sqlcipher3)
                pass  # Column already exists

        # Ensure observations has invalidated_by column (contradiction invalidation)
        try:
            self._conn.execute(
                "ALTER TABLE observations ADD COLUMN invalidated_by TEXT"
            )
        except Exception:  # Column/table may already exist (sqlite3 or sqlcipher3)
            pass  # Column already exists

        # Ensure observations has FSRS review tracking columns
        for col, col_type in [
            ("review_count", "INTEGER DEFAULT 0"),
            ("last_reviewed_at", "TEXT"),
            ("review_source", "TEXT DEFAULT ''"),
        ]:
            try:
                self._conn.execute(
                    f"ALTER TABLE observations ADD COLUMN {col} {col_type}"
                )
            except Exception:  # Column/table may already exist (sqlite3 or sqlcipher3)
                pass  # Column already exists

        # Ensure patterns has confirmations column (legacy DBs had only 'applications')
        for col, col_type in [
            ("confirmations", "INTEGER DEFAULT 0"),
        ]:
            try:
                self._conn.execute(
                    f"ALTER TABLE patterns ADD COLUMN {col} {col_type}"
                )
            except Exception:  # Column/table may already exist (sqlite3 or sqlcipher3)
                pass  # Column already exists

        # Work item review states, result summaries, and mode tracking
        for col, col_type in [
            ("review_notes", "TEXT DEFAULT NULL"),
            ("result_summary", "TEXT"),
            ("dispatch_mode", "TEXT"),
        ]:
            try:
                self._conn.execute(
                    f"ALTER TABLE work_items ADD COLUMN {col} {col_type}"
                )
            except Exception:  # Column/table may already exist (sqlite3 or sqlcipher3)
                pass  # Column already exists

        # Ensure knowledge_entries has all required columns
        for col, col_type in [
            ("entry_type", "TEXT"),
            ("tags", "TEXT"),
            ("applied_count", "INTEGER DEFAULT 0"),
            ("confidence", "REAL DEFAULT 0.5"),
            ("classification_level", "INTEGER NOT NULL DEFAULT 3"),
            ("last_validated_at", "TEXT"),
            ("validation_status", "TEXT"),
        ]:
            try:
                self._conn.execute(
                    f"ALTER TABLE knowledge_entries ADD COLUMN {col} {col_type}"
                )
            except Exception:  # Column/table may already exist (sqlite3 or sqlcipher3)
                pass  # Column already exists

        # Backfill entry_type from legacy 'type' column if entry_type is NULL
        try:
            self._conn.execute("""
                UPDATE knowledge_entries SET entry_type = type
                WHERE entry_type IS NULL AND type IS NOT NULL
            """)
        except Exception:  # Column/table may already exist (sqlite3 or sqlcipher3)
            pass  # 'type' column may not exist in fresh DBs

        # Backfill confirmations from legacy 'applications' column
        try:
            self._conn.execute("""
                UPDATE patterns SET confirmations = applications
                WHERE confirmations = 0 AND applications > 0
            """)
        except Exception:  # Column/table may already exist (sqlite3 or sqlcipher3)
            pass  # 'applications' column may not exist in fresh DBs

        # Core sqlite-vec virtual tables for semantic search
        # Wrapped in try/except: sqlite-vec shadow tables may already exist
        # from a prior migration, causing CREATE VIRTUAL TABLE to fail even
        # with IF NOT EXISTS (known sqlite-vec quirk).
        for table, pk_col in [
            ("vec_documents", "document_id"),
            ("vec_knowledge", "entry_id"),
            ("vec_observations", "observation_id"),
            ("vec_code_symbols", "symbol_id"),
        ]:
            try:
                self._conn.execute(f"""
                    CREATE VIRTUAL TABLE IF NOT EXISTS {table} USING vec0(
                        {pk_col} TEXT PRIMARY KEY,
                        embedding float[768]
                    )
                """)
            except Exception:  # Column/table may already exist (sqlite3 or sqlcipher3)
                pass  # Shadow tables already exist
        try:
            self._conn.execute("""
                CREATE VIRTUAL TABLE IF NOT EXISTS vec_chunks USING vec0(
                    chunk_id INTEGER PRIMARY KEY,
                    embedding float[768]
                )
            """)
        except Exception:  # Column/table may already exist (sqlite3 or sqlcipher3)
            pass  # Shadow tables already exist

        # Security audit log — append-only record of scanner events
        self._conn.executescript("""
            CREATE TABLE IF NOT EXISTS security_events (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                event_type TEXT NOT NULL,
                severity TEXT NOT NULL,
                pattern_name TEXT,
                tool_name TEXT,
                action_taken TEXT NOT NULL,
                detail TEXT,
                created_at TEXT DEFAULT (datetime('now'))
            );
            CREATE INDEX IF NOT EXISTS idx_security_events_type ON security_events(event_type);
            CREATE INDEX IF NOT EXISTS idx_security_events_severity ON security_events(severity);
            CREATE INDEX IF NOT EXISTS idx_security_events_created ON security_events(created_at);
        """)
        self._conn.commit()

    def execute(self, query: str, params: tuple = ()) -> list[dict[str, Any]]:
        """Execute query and return list of dicts."""
        cursor = self._conn.execute(query, params)
        if cursor.description is None:
            if not self._in_transaction:
                self._conn.commit()
            return []
        return [dict(row) for row in cursor.fetchall()]

    def execute_one(self, query: str, params: tuple = ()) -> dict[str, Any] | None:
        """Execute query and return first row as dict, or None."""
        cursor = self._conn.execute(query, params)
        if cursor.description is None:
            if not self._in_transaction:
                self._conn.commit()
            return None
        row = cursor.fetchone()
        return dict(row) if row else None

    @contextmanager
    def transaction(self):
        """Context manager for explicit transactions."""
        self._conn.execute("BEGIN")
        self._in_transaction = True
        try:
            yield self
            self._conn.commit()
        except Exception:
            self._conn.rollback()
            raise
        finally:
            self._in_transaction = False

    def close(self):
        """Close the database connection."""
        self._conn.close()


# ---------------------------------------------------------------------------
# CockroachDB / PostgreSQL backend
# ---------------------------------------------------------------------------

# SSL certificate defaults (same paths as lib/db.py for backward compat)
_SSL_ROOT_CERT = os.environ.get(
    "CRDB_SSL_ROOT_CERT", os.path.expanduser("~/.cortex/certs/ca.crt")
)
_SSL_CERT = os.environ.get(
    "CRDB_SSL_CERT", os.path.expanduser("~/.cortex/certs/client/client.root.crt")
)
_SSL_KEY = os.environ.get(
    "CRDB_SSL_KEY", os.path.expanduser("~/.cortex/certs/client/client.root.key")
)

# Regex: match single-quoted or double-quoted strings, or a bare `?`
_PLACEHOLDER_RE = re.compile(r"'[^']*'|\"[^\"]*\"|\?")


def build_connection_string() -> str:
    """Build CockroachDB connection string with SSL certificates.

    Uses GREYMATTER_DB_URL if set, otherwise constructs from CRDB_* env vars.
    """
    base_url = os.environ.get("GREYMATTER_DB_URL")
    if base_url:
        return base_url

    host = os.environ.get("CRDB_HOST", "localhost")
    port = os.environ.get("CRDB_PORT", "26257")
    user = os.environ.get("CRDB_USER", "root")
    database = os.environ.get("CRDB_DATABASE", "greymatter")

    return (
        f"postgresql://{user}@{host}:{port}/{database}"
        f"?sslmode=verify-full"
        f"&sslrootcert={_SSL_ROOT_CERT}"
        f"&sslcert={_SSL_CERT}"
        f"&sslkey={_SSL_KEY}"
    )


class CockroachDatabase:
    """CockroachDB database for distributed mode.

    Provides the same Database protocol as SQLiteDatabase but backed by
    CockroachDB via psycopg2 with connection pooling.

    Queries written with SQLite-style ``?`` placeholders are automatically
    converted to PostgreSQL ``%s`` placeholders, so the same query strings
    work against both backends.
    """

    @property
    def dialect(self) -> str:
        return "crdb"

    def __init__(
        self,
        dsn: str | None = None,
        minconn: int = 2,
        maxconn: int = 20,
    ):
        if not PSYCOPG2_AVAILABLE:
            raise ImportError(
                "psycopg2 is required for CockroachDB connections. "
                "Install with: pip install psycopg2-binary"
            )

        self._dsn = dsn or build_connection_string()
        self._pool = pg_pool.ThreadedConnectionPool(
            minconn, maxconn, self._dsn
        )
        self._local = threading.local()

    # -- placeholder conversion ------------------------------------------------

    @staticmethod
    def _convert_placeholders(query: str) -> str:
        """Convert SQLite ``?`` placeholders to PostgreSQL ``%s``.

        Preserves ``?`` characters inside single- or double-quoted strings.
        """
        if "?" not in query:
            return query

        def _replace(match: re.Match) -> str:
            token = match.group(0)
            if token == "?":
                return "%s"
            return token  # quoted string — leave unchanged

        return _PLACEHOLDER_RE.sub(_replace, query)

    # -- connection helpers ----------------------------------------------------

    def _get_conn(self):
        """Get a connection from the pool (thread-local)."""
        conn = getattr(self._local, "conn", None)
        if conn is None or conn.closed:
            conn = self._pool.getconn()
            conn.autocommit = True
            self._local.conn = conn
        return conn

    def _put_conn(self, conn):
        """Return a connection to the pool."""
        self._pool.putconn(conn)
        self._local.conn = None

    # -- Database protocol -----------------------------------------------------

    def execute(self, query: str, params: tuple = ()) -> list[dict[str, Any]]:
        """Execute query and return list of dicts."""
        pg_query = self._convert_placeholders(query)
        conn = self._get_conn()
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(pg_query, params)
            if cur.description is None:
                return []
            return [dict(row) for row in cur.fetchall()]

    def execute_one(self, query: str, params: tuple = ()) -> dict[str, Any] | None:
        """Execute query and return first row as dict, or None."""
        pg_query = self._convert_placeholders(query)
        conn = self._get_conn()
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(pg_query, params)
            if cur.description is None:
                return None
            row = cur.fetchone()
            return dict(row) if row else None

    @contextmanager
    def transaction(self):
        """Context manager for explicit transactions."""
        conn = self._get_conn()
        old_autocommit = conn.autocommit
        conn.autocommit = False
        try:
            yield self
            conn.commit()
        except Exception:
            conn.rollback()
            raise
        finally:
            conn.autocommit = old_autocommit

    def migrate(self) -> None:
        """Run schema migrations.

        CockroachDB schema is managed by the schema/migrations module.
        This is a no-op here — call the migration runner separately.
        """
        pass

    def close(self):
        """Close all connections in the pool."""
        self._pool.closeall()


# ---------------------------------------------------------------------------
# Async wrapper for any Database implementation
# ---------------------------------------------------------------------------


class AsyncDatabase:
    """Async wrapper around a synchronous Database.

    Delegates blocking execute/execute_one/transaction calls to a thread pool
    via asyncio.loop.run_in_executor, keeping the event loop responsive.
    The underlying sync Database is still available as .sync for code that
    needs direct access (e.g., schema extensions, migrations).
    """

    def __init__(self, db: Database) -> None:
        self._db = db

    @property
    def sync(self) -> Database:
        """Access the underlying synchronous Database."""
        return self._db

    @property
    def dialect(self) -> str:
        """Return the dialect of the underlying database."""
        return self._db.dialect

    async def execute(self, query: str, params: tuple = ()) -> list[dict[str, Any]]:
        """Execute query in a thread pool and return list of dicts."""
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(None, self._db.execute, query, params)

    async def execute_one(self, query: str, params: tuple = ()) -> dict[str, Any] | None:
        """Execute query in a thread pool and return first row or None."""
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(None, self._db.execute_one, query, params)

    @asynccontextmanager
    async def transaction(self):
        """Async context manager for transactions.

        Wraps the sync database's transaction() in run_in_executor calls.
        Works with both SQLiteDatabase and CockroachDatabase.
        """
        loop = asyncio.get_running_loop()
        ctx = self._db.transaction()
        await loop.run_in_executor(None, ctx.__enter__)
        error = None
        try:
            yield self
        except BaseException as exc:
            error = exc
        # Always exit the sync context manager in the executor
        try:
            if error is None:
                await loop.run_in_executor(None, ctx.__exit__, None, None, None)
            else:
                suppressed = await loop.run_in_executor(
                    None, ctx.__exit__, type(error), error, error.__traceback__
                )
                if not suppressed:
                    raise error
        except BaseException:
            if error is not None:
                raise error
            raise

    async def close(self) -> None:
        """Close the underlying database connection."""
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(None, self._db.close)
