"""
database.py — PostgreSQL database with asyncpg connection pool.

Application-level AES-256-GCM encryption for credentials (unchanged from SQLite era).
The pool is initialised once at startup via init_db() and closed via close_db().
All store methods are async — callers must await them.
"""
from __future__ import annotations

import base64
import json
import os
from datetime import datetime, timezone
from typing import Any
from urllib.parse import urlparse

import asyncpg
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

from .config import settings

# ─── Encryption ───────────────────────────────────────────────────────────────
# Unchanged from SQLite version — encrypted blobs are stored as base64 TEXT.

_SALT = b"aide-credential-store-v1"
_ITERATIONS = 480_000


def _derive_key(password: str) -> bytes:
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=_SALT,
        iterations=_ITERATIONS,
    )
    return kdf.derive(password.encode())


_ENCRYPTION_KEY = _derive_key(settings.db_master_password)


def _encrypt(plaintext: str) -> str:
    """Encrypt a string value, return base64-encoded ciphertext (nonce + tag + data)."""
    aesgcm = AESGCM(_ENCRYPTION_KEY)
    nonce = os.urandom(12)
    ciphertext = aesgcm.encrypt(nonce, plaintext.encode(), None)
    return base64.b64encode(nonce + ciphertext).decode()


def _decrypt(encoded: str) -> str:
    """Decrypt a base64-encoded ciphertext, return plaintext string."""
    data = base64.b64decode(encoded)
    nonce, ciphertext = data[:12], data[12:]
    aesgcm = AESGCM(_ENCRYPTION_KEY)
    return aesgcm.decrypt(nonce, ciphertext, None).decode()


# ─── Connection Pool ──────────────────────────────────────────────────────────

_pool: asyncpg.Pool | None = None


async def get_pool() -> asyncpg.Pool:
    """Return the shared connection pool. Must call init_db() first."""
    assert _pool is not None, "Database not initialised — call init_db() first"
    return _pool


# ─── Migrations ───────────────────────────────────────────────────────────────
# Each migration is a list of SQL statements (asyncpg runs one statement at a time).
# All migrations are idempotent (IF NOT EXISTS / ADD COLUMN IF NOT EXISTS / ON CONFLICT DO NOTHING).

_MIGRATIONS: list[list[str]] = [
    # v1 — initial schema
    [
        """CREATE TABLE IF NOT EXISTS schema_version (
            version INTEGER PRIMARY KEY
        )""",
        """CREATE TABLE IF NOT EXISTS credentials (
            key         TEXT PRIMARY KEY,
            value_enc   TEXT NOT NULL,
            description TEXT,
            created_at  TEXT NOT NULL,
            updated_at  TEXT NOT NULL
        )""",
        """CREATE TABLE IF NOT EXISTS audit_log (
            id             BIGSERIAL PRIMARY KEY,
            timestamp      TEXT NOT NULL,
            session_id     TEXT,
            tool_name      TEXT NOT NULL,
            arguments      JSONB,
            result_summary TEXT,
            confirmed      BOOLEAN NOT NULL DEFAULT FALSE,
            task_id        TEXT
        )""",
        "CREATE INDEX IF NOT EXISTS idx_audit_timestamp ON audit_log(timestamp)",
        "CREATE INDEX IF NOT EXISTS idx_audit_session   ON audit_log(session_id)",
        "CREATE INDEX IF NOT EXISTS idx_audit_tool      ON audit_log(tool_name)",
        """CREATE TABLE IF NOT EXISTS scheduled_tasks (
            id            TEXT PRIMARY KEY,
            name          TEXT NOT NULL,
            description   TEXT,
            schedule      TEXT,
            prompt        TEXT NOT NULL,
            allowed_tools JSONB,
            enabled       BOOLEAN NOT NULL DEFAULT TRUE,
            last_run      TEXT,
            last_status   TEXT,
            created_at    TEXT NOT NULL,
            updated_at    TEXT NOT NULL
        )""",
        """CREATE TABLE IF NOT EXISTS conversations (
            id         TEXT PRIMARY KEY,
            started_at TEXT NOT NULL,
            ended_at   TEXT,
            messages   JSONB NOT NULL,
            task_id    TEXT
        )""",
    ],
    # v2 — email whitelist, agents, agent_runs
    [
        """CREATE TABLE IF NOT EXISTS email_whitelist (
            email       TEXT PRIMARY KEY,
            daily_limit INTEGER NOT NULL DEFAULT 0,
            created_at  TEXT NOT NULL
        )""",
        """CREATE TABLE IF NOT EXISTS agents (
            id                   TEXT PRIMARY KEY,
            name                 TEXT NOT NULL,
            description          TEXT,
            prompt               TEXT NOT NULL,
            model                TEXT NOT NULL,
            can_create_subagents BOOLEAN NOT NULL DEFAULT FALSE,
            allowed_tools        JSONB,
            schedule             TEXT,
            enabled              BOOLEAN NOT NULL DEFAULT TRUE,
            parent_agent_id      TEXT REFERENCES agents(id),
            created_by           TEXT NOT NULL DEFAULT 'user',
            created_at           TEXT NOT NULL,
            updated_at           TEXT NOT NULL
        )""",
        """CREATE TABLE IF NOT EXISTS agent_runs (
            id            TEXT PRIMARY KEY,
            agent_id      TEXT NOT NULL REFERENCES agents(id),
            started_at    TEXT NOT NULL,
            ended_at      TEXT,
            status        TEXT NOT NULL DEFAULT 'running',
            input_tokens  INTEGER NOT NULL DEFAULT 0,
            output_tokens INTEGER NOT NULL DEFAULT 0,
            cost_usd      REAL,
            result        TEXT,
            error         TEXT
        )""",
        "CREATE INDEX IF NOT EXISTS idx_agent_runs_agent_id   ON agent_runs(agent_id)",
        "CREATE INDEX IF NOT EXISTS idx_agent_runs_started_at ON agent_runs(started_at)",
        "CREATE INDEX IF NOT EXISTS idx_agent_runs_status     ON agent_runs(status)",
    ],
    # v3 — web domain whitelist
    [
        """CREATE TABLE IF NOT EXISTS web_whitelist (
            domain     TEXT PRIMARY KEY,
            note       TEXT NOT NULL DEFAULT '',
            created_at TEXT NOT NULL
        )""",
        "INSERT INTO web_whitelist (domain, note, created_at) VALUES ('duckduckgo.com', 'DuckDuckGo search', '2024-01-01T00:00:00+00:00') ON CONFLICT DO NOTHING",
        "INSERT INTO web_whitelist (domain, note, created_at) VALUES ('wikipedia.org', 'Wikipedia', '2024-01-01T00:00:00+00:00') ON CONFLICT DO NOTHING",
        "INSERT INTO web_whitelist (domain, note, created_at) VALUES ('weather.met.no', 'Norwegian Meteorological Institute', '2024-01-01T00:00:00+00:00') ON CONFLICT DO NOTHING",
        "INSERT INTO web_whitelist (domain, note, created_at) VALUES ('api.met.no', 'Norwegian Meteorological API', '2024-01-01T00:00:00+00:00') ON CONFLICT DO NOTHING",
        "INSERT INTO web_whitelist (domain, note, created_at) VALUES ('yr.no', 'Yr weather service', '2024-01-01T00:00:00+00:00') ON CONFLICT DO NOTHING",
        "INSERT INTO web_whitelist (domain, note, created_at) VALUES ('timeanddate.com', 'Time and Date', '2024-01-01T00:00:00+00:00') ON CONFLICT DO NOTHING",
    ],
    # v4 — filesystem sandbox whitelist
    [
        """CREATE TABLE IF NOT EXISTS filesystem_whitelist (
            path       TEXT PRIMARY KEY,
            note       TEXT NOT NULL DEFAULT '',
            created_at TEXT NOT NULL
        )""",
    ],
    # v5 — optional agent assignment for scheduled tasks
    [
        "ALTER TABLE scheduled_tasks ADD COLUMN IF NOT EXISTS agent_id TEXT REFERENCES agents(id)",
    ],
    # v6 — per-agent max_tool_calls override
    [
        "ALTER TABLE agents ADD COLUMN IF NOT EXISTS max_tool_calls INTEGER",
    ],
    # v7 — email inbox trigger rules
    [
        """CREATE TABLE IF NOT EXISTS email_triggers (
            id           TEXT PRIMARY KEY,
            trigger_word TEXT NOT NULL,
            agent_id     TEXT NOT NULL,
            description  TEXT NOT NULL DEFAULT '',
            enabled      BOOLEAN NOT NULL DEFAULT TRUE,
            created_at   TEXT NOT NULL,
            updated_at   TEXT NOT NULL
        )""",
    ],
    # v8 — Telegram bot integration
    [
        """CREATE TABLE IF NOT EXISTS telegram_whitelist (
            chat_id    TEXT PRIMARY KEY,
            label      TEXT NOT NULL DEFAULT '',
            created_at TEXT NOT NULL
        )""",
        """CREATE TABLE IF NOT EXISTS telegram_triggers (
            id           TEXT PRIMARY KEY,
            trigger_word TEXT NOT NULL,
            agent_id     TEXT NOT NULL,
            description  TEXT NOT NULL DEFAULT '',
            enabled      BOOLEAN NOT NULL DEFAULT TRUE,
            created_at   TEXT NOT NULL,
            updated_at   TEXT NOT NULL
        )""",
    ],
    # v9 — agent prompt_mode column
    [
        "ALTER TABLE agents ADD COLUMN IF NOT EXISTS prompt_mode TEXT NOT NULL DEFAULT 'combined'",
    ],
    # v10 — (was SQLite re-apply of v9; no-op here)
    [],
    # v11 — MCP client server configurations
    [
        """CREATE TABLE IF NOT EXISTS mcp_servers (
            id          TEXT PRIMARY KEY,
            name        TEXT NOT NULL,
            url         TEXT NOT NULL,
            transport   TEXT NOT NULL DEFAULT 'sse',
            api_key_enc TEXT,
            headers_enc TEXT,
            enabled     BOOLEAN NOT NULL DEFAULT TRUE,
            created_at  TEXT NOT NULL,
            updated_at  TEXT NOT NULL
        )""",
    ],
    # v12 — users table for multi-user support (Part 2)
    [
        """CREATE TABLE IF NOT EXISTS users (
            id            TEXT PRIMARY KEY,
            username      TEXT NOT NULL UNIQUE,
            password_hash TEXT NOT NULL,
            role          TEXT NOT NULL DEFAULT 'user',
            is_active     BOOLEAN NOT NULL DEFAULT TRUE,
            totp_secret   TEXT,
            created_at    TEXT NOT NULL,
            updated_at    TEXT NOT NULL
        )""",
        "CREATE INDEX IF NOT EXISTS idx_users_username ON users(username)",
        "ALTER TABLE agents      ADD COLUMN IF NOT EXISTS owner_user_id TEXT REFERENCES users(id)",
        "ALTER TABLE conversations ADD COLUMN IF NOT EXISTS user_id TEXT REFERENCES users(id)",
        "ALTER TABLE audit_log   ADD COLUMN IF NOT EXISTS user_id TEXT REFERENCES users(id)",
    ],
    # v13 — add email column to users
    [
        "ALTER TABLE users ADD COLUMN IF NOT EXISTS email TEXT",
    ],
    # v14 — per-user settings table + user_id columns on multi-tenant tables
    [
        """CREATE TABLE IF NOT EXISTS user_settings (
            user_id  TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
            key      TEXT NOT NULL,
            value    TEXT,
            PRIMARY KEY (user_id, key)
        )""",
        "ALTER TABLE email_triggers    ADD COLUMN IF NOT EXISTS user_id TEXT REFERENCES users(id)",
        "ALTER TABLE telegram_triggers ADD COLUMN IF NOT EXISTS user_id TEXT REFERENCES users(id)",
        "ALTER TABLE telegram_whitelist ADD COLUMN IF NOT EXISTS user_id TEXT REFERENCES users(id)",
        "ALTER TABLE mcp_servers       ADD COLUMN IF NOT EXISTS user_id TEXT REFERENCES users(id)",
    ],
    # v15 — fix telegram_whitelist unique constraint to allow (chat_id, user_id) pairs
    # Uses NULLS NOT DISTINCT (PostgreSQL 15+) so (chat_id, NULL) is unique per global entry
    [
        # Drop old primary key constraint so chat_id alone no longer enforces uniqueness
        """DO $$ BEGIN
            IF EXISTS (
                SELECT 1 FROM pg_constraint
                WHERE conname = 'telegram_whitelist_pkey' AND conrelid = 'telegram_whitelist'::regclass
            ) THEN
                ALTER TABLE telegram_whitelist DROP CONSTRAINT telegram_whitelist_pkey;
            END IF;
        END $$""",
        # Add a surrogate UUID primary key
        "ALTER TABLE telegram_whitelist ADD COLUMN IF NOT EXISTS id UUID DEFAULT gen_random_uuid()",
        # Make it not null and set primary key (only if not already set)
        """DO $$ BEGIN
            IF NOT EXISTS (
                SELECT 1 FROM pg_constraint
                WHERE conname = 'telegram_whitelist_pk' AND conrelid = 'telegram_whitelist'::regclass
            ) THEN
                ALTER TABLE telegram_whitelist ADD CONSTRAINT telegram_whitelist_pk PRIMARY KEY (id);
            END IF;
        END $$""",
        # Create unique index on (chat_id, user_id) NULLS NOT DISTINCT
        """CREATE UNIQUE INDEX IF NOT EXISTS telegram_whitelist_chat_user_idx
           ON telegram_whitelist (chat_id, user_id) NULLS NOT DISTINCT""",
    ],
    # v16 — email_accounts table for multi-account email handling
    [
        """CREATE TABLE IF NOT EXISTS email_accounts (
            id                 UUID PRIMARY KEY DEFAULT gen_random_uuid(),
            user_id            TEXT REFERENCES users(id),
            label              TEXT NOT NULL,
            account_type       TEXT NOT NULL DEFAULT 'handling',
            imap_host          TEXT NOT NULL,
            imap_port          INTEGER NOT NULL DEFAULT 993,
            imap_username      TEXT NOT NULL,
            imap_password      TEXT NOT NULL,
            smtp_host          TEXT,
            smtp_port          INTEGER,
            smtp_username      TEXT,
            smtp_password      TEXT,
            agent_id           TEXT REFERENCES agents(id),
            enabled            BOOLEAN NOT NULL DEFAULT TRUE,
            initial_load_done  BOOLEAN NOT NULL DEFAULT FALSE,
            initial_load_limit INTEGER NOT NULL DEFAULT 200,
            monitored_folders  TEXT NOT NULL DEFAULT '[\"INBOX\"]',
            created_at         TEXT NOT NULL,
            updated_at         TEXT NOT NULL
        )""",
        "ALTER TABLE email_triggers ADD COLUMN IF NOT EXISTS account_id UUID REFERENCES email_accounts(id)",
    ],
    # v17 — convert audit_log.arguments from TEXT to JSONB (SQLite-migrated DBs have TEXT)
    #        and agents/scheduled_tasks allowed_tools from TEXT to JSONB if not already
    [
        """DO $$
        BEGIN
            IF (SELECT data_type FROM information_schema.columns
                WHERE table_name='audit_log' AND column_name='arguments') = 'text' THEN
                ALTER TABLE audit_log
                    ALTER COLUMN arguments TYPE JSONB
                    USING CASE WHEN arguments IS NULL OR arguments = '' THEN NULL
                               ELSE arguments::jsonb END;
            END IF;
        END $$""",
        """DO $$
        BEGIN
            IF (SELECT data_type FROM information_schema.columns
                WHERE table_name='agents' AND column_name='allowed_tools') = 'text' THEN
                ALTER TABLE agents
                    ALTER COLUMN allowed_tools TYPE JSONB
                    USING CASE WHEN allowed_tools IS NULL OR allowed_tools = '' THEN NULL
                               ELSE allowed_tools::jsonb END;
            END IF;
        END $$""",
        """DO $$
        BEGIN
            IF (SELECT data_type FROM information_schema.columns
                WHERE table_name='scheduled_tasks' AND column_name='allowed_tools') = 'text' THEN
                ALTER TABLE scheduled_tasks
                    ALTER COLUMN allowed_tools TYPE JSONB
                    USING CASE WHEN allowed_tools IS NULL OR allowed_tools = '' THEN NULL
                               ELSE allowed_tools::jsonb END;
            END IF;
        END $$""",
    ],
    # v18 — MFA challenge table for TOTP second-factor login
    [
        """CREATE TABLE IF NOT EXISTS mfa_challenges (
            token      TEXT PRIMARY KEY,
            user_id    TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
            next_url   TEXT NOT NULL DEFAULT '/',
            created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
            expires_at TIMESTAMPTZ NOT NULL,
            attempts   INTEGER NOT NULL DEFAULT 0
        )""",
        "CREATE INDEX IF NOT EXISTS idx_mfa_challenges_expires ON mfa_challenges(expires_at)",
    ],
    # v19 — display name for users (editable, separate from username)
    [
        "ALTER TABLE users ADD COLUMN IF NOT EXISTS display_name TEXT",
    ],
    # v20 — extra notification tools for handling email accounts
    [
        "ALTER TABLE email_accounts ADD COLUMN IF NOT EXISTS extra_tools JSONB DEFAULT '[]'",
    ],
    # v21 — bound Telegram chat_id for email handling accounts
    [
        "ALTER TABLE email_accounts ADD COLUMN IF NOT EXISTS telegram_chat_id TEXT",
    ],
    # v22 — Telegram keyword routing + pause flag for email handling accounts
    [
        "ALTER TABLE email_accounts ADD COLUMN IF NOT EXISTS telegram_keyword TEXT",
        "ALTER TABLE email_accounts ADD COLUMN IF NOT EXISTS paused BOOLEAN DEFAULT FALSE",
    ],
    # v23 — Conversation title for chat history UI
    [
        "ALTER TABLE conversations ADD COLUMN IF NOT EXISTS title TEXT",
    ],
    # v24 — Store model ID used in each conversation
    [
        "ALTER TABLE conversations ADD COLUMN IF NOT EXISTS model TEXT",
    ],
    # v25 — Inbound webhook trigger endpoints + outbound webhook targets + monitors
    [
        """
        CREATE TABLE IF NOT EXISTS webhook_endpoints (
            id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
            name TEXT NOT NULL,
            token TEXT NOT NULL UNIQUE,
            description TEXT DEFAULT '',
            agent_id TEXT REFERENCES agents(id) ON DELETE CASCADE,
            enabled BOOLEAN DEFAULT TRUE,
            allow_get BOOLEAN DEFAULT TRUE,
            created_at TIMESTAMPTZ DEFAULT NOW(),
            last_triggered_at TIMESTAMPTZ,
            trigger_count INTEGER DEFAULT 0
        )
        """,
        """
        CREATE TABLE IF NOT EXISTS webhook_targets (
            id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
            name TEXT NOT NULL UNIQUE,
            url TEXT NOT NULL,
            secret_header TEXT,
            enabled BOOLEAN DEFAULT TRUE,
            created_at TIMESTAMPTZ DEFAULT NOW()
        )
        """,
        """
        CREATE TABLE IF NOT EXISTS watched_pages (
            id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
            name TEXT NOT NULL,
            url TEXT NOT NULL,
            schedule TEXT NOT NULL DEFAULT '0 * * * *',
            css_selector TEXT,
            last_content_hash TEXT,
            last_changed_at TIMESTAMPTZ,
            last_checked_at TIMESTAMPTZ,
            last_error TEXT,
            agent_id TEXT REFERENCES agents(id) ON DELETE SET NULL,
            notification_mode TEXT DEFAULT 'agent',
            enabled BOOLEAN DEFAULT TRUE,
            owner_user_id TEXT REFERENCES users(id) ON DELETE CASCADE,
            created_at TIMESTAMPTZ DEFAULT NOW()
        )
        """,
        """
        CREATE TABLE IF NOT EXISTS rss_feeds (
            id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
            name TEXT NOT NULL,
            url TEXT NOT NULL,
            schedule TEXT NOT NULL DEFAULT '0 */4 * * *',
            agent_id TEXT REFERENCES agents(id) ON DELETE SET NULL,
            last_fetched_at TIMESTAMPTZ,
            last_etag TEXT,
            last_modified TEXT,
            seen_item_ids JSONB DEFAULT '[]',
            max_items_per_run INTEGER DEFAULT 5,
            notification_mode TEXT DEFAULT 'agent',
            last_error TEXT,
            enabled BOOLEAN DEFAULT TRUE,
            owner_user_id TEXT REFERENCES users(id) ON DELETE CASCADE,
            created_at TIMESTAMPTZ DEFAULT NOW()
        )
        """,
    ],
    # v26 — Per-user webhook endpoint ownership
    [
        "ALTER TABLE webhook_endpoints ADD COLUMN IF NOT EXISTS owner_user_id TEXT REFERENCES users(id) ON DELETE CASCADE",
    ],
    # v27 — Per-user webhook target ownership
    [
        "ALTER TABLE webhook_targets ADD COLUMN IF NOT EXISTS owner_user_id TEXT REFERENCES users(id) ON DELETE CASCADE",
    ],
    # v28 — Browser trusted domains (per-user interaction pre-approval)
    [
        """
        CREATE TABLE IF NOT EXISTS browser_approved_domains (
            id          UUID PRIMARY KEY DEFAULT gen_random_uuid(),
            owner_user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
            domain      TEXT NOT NULL,
            note        TEXT,
            created_at  TIMESTAMPTZ DEFAULT NOW(),
            UNIQUE (owner_user_id, domain)
        )
        """,
    ],
    # v29 — Track model per agent run for usage/cost overview
    [
        "ALTER TABLE agent_runs ADD COLUMN IF NOT EXISTS model TEXT",
    ],
    # v30 — Track token usage and cost per chat conversation
    [
        "ALTER TABLE conversations ADD COLUMN IF NOT EXISTS input_tokens  INTEGER NOT NULL DEFAULT 0",
        "ALTER TABLE conversations ADD COLUMN IF NOT EXISTS output_tokens INTEGER NOT NULL DEFAULT 0",
        "ALTER TABLE conversations ADD COLUMN IF NOT EXISTS cost_usd      REAL",
    ],
]


async def _run_migrations(conn: asyncpg.Connection) -> None:
    """Apply pending migrations idempotently, each in its own transaction."""
    await conn.execute(
        "CREATE TABLE IF NOT EXISTS schema_version (version INTEGER PRIMARY KEY)"
    )
    current: int = await conn.fetchval(
        "SELECT COALESCE(MAX(version), 0) FROM schema_version"
    ) or 0

    for i, statements in enumerate(_MIGRATIONS, start=1):
        if i <= current:
            continue
        async with conn.transaction():
            for sql in statements:
                sql = sql.strip()
                if sql:
                    await conn.execute(sql)
            await conn.execute(
                "INSERT INTO schema_version (version) VALUES ($1) ON CONFLICT DO NOTHING", i
            )
        print(f"[aide] Applied database migration v{i}")


# ─── Helpers ──────────────────────────────────────────────────────────────────

def _utcnow() -> str:
    return datetime.now(timezone.utc).isoformat()


def _jsonify(obj: Any) -> Any:
    """Return a JSON-safe version of obj (converts non-serializable values to strings)."""
    if obj is None:
        return None
    return json.loads(json.dumps(obj, default=str))


def _rowcount(status: str) -> int:
    """Parse asyncpg execute() status string like 'DELETE 3' → 3."""
    try:
        return int(status.split()[-1])
    except (ValueError, IndexError):
        return 0


# ─── Credential Store ─────────────────────────────────────────────────────────

class CredentialStore:
    """Encrypted key-value store for sensitive credentials."""

    async def get(self, key: str) -> str | None:
        pool = await get_pool()
        row = await pool.fetchrow(
            "SELECT value_enc FROM credentials WHERE key = $1", key
        )
        if row is None:
            return None
        return _decrypt(row["value_enc"])

    async def set(self, key: str, value: str, description: str = "") -> None:
        now = _utcnow()
        encrypted = _encrypt(value)
        pool = await get_pool()
        await pool.execute(
            """
            INSERT INTO credentials (key, value_enc, description, created_at, updated_at)
            VALUES ($1, $2, $3, $4, $5)
            ON CONFLICT (key) DO UPDATE SET
                value_enc   = EXCLUDED.value_enc,
                description = EXCLUDED.description,
                updated_at  = EXCLUDED.updated_at
            """,
            key, encrypted, description, now, now,
        )

    async def delete(self, key: str) -> bool:
        pool = await get_pool()
        status = await pool.execute("DELETE FROM credentials WHERE key = $1", key)
        return _rowcount(status) > 0

    async def list_keys(self) -> list[dict]:
        pool = await get_pool()
        rows = await pool.fetch(
            "SELECT key, description, created_at, updated_at FROM credentials ORDER BY key"
        )
        return [dict(r) for r in rows]

    async def require(self, key: str) -> str:
        value = await self.get(key)
        if not value:
            raise RuntimeError(
                f"Credential '{key}' is not configured. Add it via /settings."
            )
        return value


# Module-level singleton
credential_store = CredentialStore()


# ─── User Settings Store ──────────────────────────────────────────────────────

class UserSettingsStore:
    """Per-user key/value settings. Values are plaintext (not encrypted)."""

    async def get(self, user_id: str, key: str) -> str | None:
        pool = await get_pool()
        return await pool.fetchval(
            "SELECT value FROM user_settings WHERE user_id = $1 AND key = $2",
            user_id, key,
        )

    async def set(self, user_id: str, key: str, value: str) -> None:
        pool = await get_pool()
        await pool.execute(
            """
            INSERT INTO user_settings (user_id, key, value)
            VALUES ($1, $2, $3)
            ON CONFLICT (user_id, key) DO UPDATE SET value = EXCLUDED.value
            """,
            user_id, key, value,
        )

    async def delete(self, user_id: str, key: str) -> bool:
        pool = await get_pool()
        status = await pool.execute(
            "DELETE FROM user_settings WHERE user_id = $1 AND key = $2", user_id, key
        )
        return _rowcount(status) > 0

    async def get_with_global_fallback(self, user_id: str, key: str, global_key: str) -> str | None:
        """Try user-specific setting, fall back to global credential_store key."""
        val = await self.get(user_id, key)
        if val:
            return val
        return await credential_store.get(global_key)


# Module-level singleton
user_settings_store = UserSettingsStore()


# ─── Email Whitelist Store ────────────────────────────────────────────────────

class EmailWhitelistStore:
    """Manage allowed email recipients with optional per-address daily rate limits."""

    async def list(self) -> list[dict]:
        pool = await get_pool()
        rows = await pool.fetch(
            "SELECT email, daily_limit, created_at FROM email_whitelist ORDER BY email"
        )
        return [dict(r) for r in rows]

    async def add(self, email: str, daily_limit: int = 0) -> None:
        now = _utcnow()
        normalized = email.strip().lower()
        pool = await get_pool()
        await pool.execute(
            """
            INSERT INTO email_whitelist (email, daily_limit, created_at)
            VALUES ($1, $2, $3)
            ON CONFLICT (email) DO UPDATE SET daily_limit = EXCLUDED.daily_limit
            """,
            normalized, daily_limit, now,
        )

    async def remove(self, email: str) -> bool:
        normalized = email.strip().lower()
        pool = await get_pool()
        status = await pool.execute(
            "DELETE FROM email_whitelist WHERE email = $1", normalized
        )
        return _rowcount(status) > 0

    async def get(self, email: str) -> dict | None:
        normalized = email.strip().lower()
        pool = await get_pool()
        row = await pool.fetchrow(
            "SELECT email, daily_limit, created_at FROM email_whitelist WHERE email = $1",
            normalized,
        )
        return dict(row) if row else None

    async def check_rate_limit(self, email: str) -> tuple[bool, int, int]:
        """
        Check whether sending to this address is within the daily limit.
        Returns (allowed, count_today, limit).  limit=0 means unlimited.
        """
        entry = await self.get(email)
        if entry is None:
            return False, 0, 0

        limit = entry["daily_limit"]
        if limit == 0:
            return True, 0, 0

        # Compute start of today in UTC as ISO8601 string for TEXT comparison
        today_start = (
            datetime.now(timezone.utc)
            .replace(hour=0, minute=0, second=0, microsecond=0)
            .isoformat()
        )
        pool = await get_pool()
        count: int = await pool.fetchval(
            """
            SELECT COUNT(*) FROM audit_log
            WHERE tool_name = 'email'
              AND arguments->>'operation' = 'send_email'
              AND arguments->>'to' = $1
              AND timestamp >= $2
              AND (result_summary IS NULL OR result_summary NOT LIKE '%"success": false%')
            """,
            email.strip().lower(),
            today_start,
        ) or 0

        return count < limit, count, limit


# Module-level singleton
email_whitelist_store = EmailWhitelistStore()


# ─── Web Whitelist Store ──────────────────────────────────────────────────────

class WebWhitelistStore:
    """Manage Tier-1 always-allowed web domains."""

    async def list(self) -> list[dict]:
        pool = await get_pool()
        rows = await pool.fetch(
            "SELECT domain, note, created_at FROM web_whitelist ORDER BY domain"
        )
        return [dict(r) for r in rows]

    async def add(self, domain: str, note: str = "") -> None:
        normalized = _normalize_domain(domain)
        now = _utcnow()
        pool = await get_pool()
        await pool.execute(
            """
            INSERT INTO web_whitelist (domain, note, created_at)
            VALUES ($1, $2, $3)
            ON CONFLICT (domain) DO UPDATE SET note = EXCLUDED.note
            """,
            normalized, note, now,
        )

    async def remove(self, domain: str) -> bool:
        normalized = _normalize_domain(domain)
        pool = await get_pool()
        status = await pool.execute(
            "DELETE FROM web_whitelist WHERE domain = $1", normalized
        )
        return _rowcount(status) > 0

    async def is_allowed(self, url: str) -> bool:
        """Return True if the URL's hostname matches a whitelisted domain or subdomain."""
        try:
            hostname = urlparse(url).hostname or ""
        except Exception:
            return False
        if not hostname:
            return False
        domains = await self.list()
        for entry in domains:
            d = entry["domain"]
            if hostname == d or hostname.endswith("." + d):
                return True
        return False


def _normalize_domain(domain: str) -> str:
    """Strip scheme and path, return lowercase hostname only."""
    d = domain.strip().lower()
    if "://" not in d:
        d = "https://" + d
    parsed = urlparse(d)
    return parsed.hostname or domain.strip().lower()


# Module-level singleton
web_whitelist_store = WebWhitelistStore()


# ─── Filesystem Whitelist Store ───────────────────────────────────────────────

class FilesystemWhitelistStore:
    """Manage allowed filesystem sandbox directories."""

    async def list(self) -> list[dict]:
        pool = await get_pool()
        rows = await pool.fetch(
            "SELECT path, note, created_at FROM filesystem_whitelist ORDER BY path"
        )
        return [dict(r) for r in rows]

    async def add(self, path: str, note: str = "") -> None:
        from pathlib import Path as _Path
        normalized = str(_Path(path).resolve())
        now = _utcnow()
        pool = await get_pool()
        await pool.execute(
            """
            INSERT INTO filesystem_whitelist (path, note, created_at)
            VALUES ($1, $2, $3)
            ON CONFLICT (path) DO UPDATE SET note = EXCLUDED.note
            """,
            normalized, note, now,
        )

    async def remove(self, path: str) -> bool:
        from pathlib import Path as _Path
        normalized = str(_Path(path).resolve())
        pool = await get_pool()
        status = await pool.execute(
            "DELETE FROM filesystem_whitelist WHERE path = $1", normalized
        )
        if _rowcount(status) == 0:
            # Fallback: try exact match without resolving
            status = await pool.execute(
                "DELETE FROM filesystem_whitelist WHERE path = $1", path
            )
        return _rowcount(status) > 0

    async def is_allowed(self, path: Any) -> tuple[bool, str]:
        """
        Check if path is inside any whitelisted directory.
        Returns (allowed, resolved_path_str).
        """
        from pathlib import Path as _Path
        try:
            resolved = _Path(path).resolve()
        except Exception as e:
            raise ValueError(f"Invalid path: {e}")

        sandboxes = await self.list()
        for entry in sandboxes:
            try:
                resolved.relative_to(_Path(entry["path"]).resolve())
                return True, str(resolved)
            except ValueError:
                continue
        return False, str(resolved)


# Module-level singleton
filesystem_whitelist_store = FilesystemWhitelistStore()


# ─── Initialisation ───────────────────────────────────────────────────────────

async def _init_connection(conn: asyncpg.Connection) -> None:
    """
    Register codecs on every new connection so asyncpg handles JSONB ↔ dict automatically.

    Critical: asyncpg does NOT auto-serialize Python dicts to PostgreSQL JSONB.
    Without this, inserting a dict into a JSONB column raises:
        asyncpg.exceptions.UnsupportedClientFeatureError: ...
    The codecs must be registered on every connection (not just once on the pool),
    which is why this is passed as `init=` to create_pool().
    """
    await conn.set_type_codec(
        "jsonb",
        encoder=json.dumps,
        decoder=json.loads,
        schema="pg_catalog",
    )
    await conn.set_type_codec(
        "json",
        encoder=json.dumps,
        decoder=json.loads,
        schema="pg_catalog",
    )


async def init_db() -> None:
    """Initialise the connection pool and run migrations. Call once at startup."""
    global _pool
    _pool = await asyncpg.create_pool(
        settings.aide_db_url,
        min_size=2,
        max_size=10,
        init=_init_connection,
    )
    async with _pool.acquire() as conn:
        await _run_migrations(conn)
    print(f"[aide] Database ready: {settings.aide_db_url.split('@')[-1]}")


async def close_db() -> None:
    """Close the connection pool. Call at shutdown."""
    global _pool
    if _pool:
        await _pool.close()
        _pool = None
