"""
Postgres backend — for production and enterprise deployments.

Requires psycopg2-binary:
    pip install autopil[postgres]

DATABASE_URL format:
    postgresql://user:password@host:5432/dbname
"""

import hashlib
import secrets
from contextlib import contextmanager
from datetime import datetime

try:
    import psycopg2
    import psycopg2.extras
    import psycopg2.pool
except ImportError as e:
    raise ImportError(
        "Postgres backend requires psycopg2. Install it with: pip install autopil[postgres]"
    ) from e

from autopil.models import AuditEvent, CatalogClassification, CatalogConnection, Decision, SyncResult


def _compute_pg_chain_hash(prev_hash: str, event: AuditEvent, tenant_id: str) -> str:
    """SHA-256 over the immutable audit fields and the previous hash."""
    ts = event.timestamp.isoformat() if isinstance(event.timestamp, datetime) else str(event.timestamp)
    payload = "|".join([
        prev_hash,
        event.event_id,
        tenant_id,
        ts,
        event.agent_role,
        event.user_id,
        event.source_id,
        event.decision.value,
        event.query,
    ])
    return hashlib.sha256(payload.encode()).hexdigest()

from .base import AuditLogBase, AlertStoreBase, CatalogStoreBase, KeyStoreBase, TenantStoreBase, PolicyStoreBase, TrialStoreBase

KEY_PREFIX = "apl_"


def _hash(key: str) -> str:
    return hashlib.sha256(key.encode()).hexdigest()


class _PostgresBase:
    def __init__(self, database_url: str):
        self._pool = psycopg2.pool.ThreadedConnectionPool(1, 10, dsn=database_url)

    @contextmanager
    def _conn(self):
        conn = self._pool.getconn()
        try:
            yield conn
            conn.commit()
        except Exception:
            conn.rollback()
            raise
        finally:
            self._pool.putconn(conn)


class PostgresAuditLog(_PostgresBase, AuditLogBase):

    def __init__(self, database_url: str):
        super().__init__(database_url)
        self._init_db()

    def _init_db(self):
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS audit_events (
                        event_id     TEXT PRIMARY KEY,
                        session_id   TEXT NOT NULL,
                        agent_role   TEXT NOT NULL,
                        user_id      TEXT NOT NULL,
                        source_id    TEXT NOT NULL,
                        query        TEXT NOT NULL,
                        decision     TEXT NOT NULL,
                        policy_name  TEXT NOT NULL,
                        reason       TEXT NOT NULL,
                        context_hash TEXT,
                        timestamp    TIMESTAMPTZ NOT NULL,
                        tenant_id    TEXT NOT NULL DEFAULT 'default'
                    )
                """)
                # Migrations: add columns for databases created before these fields existed
                cur.execute("ALTER TABLE audit_events ADD COLUMN IF NOT EXISTS tenant_id TEXT NOT NULL DEFAULT 'default'")
                cur.execute("ALTER TABLE audit_events ADD COLUMN IF NOT EXISTS sensitivity_level TEXT")
                cur.execute("ALTER TABLE audit_events ADD COLUMN IF NOT EXISTS source_type TEXT NOT NULL DEFAULT 'rest'")
                cur.execute("ALTER TABLE audit_events ADD COLUMN IF NOT EXISTS catalog_enriched BOOLEAN NOT NULL DEFAULT FALSE")
                cur.execute("ALTER TABLE audit_events ADD COLUMN IF NOT EXISTS catalog_connection_id TEXT")
                cur.execute("ALTER TABLE audit_events ADD COLUMN IF NOT EXISTS chain_hash TEXT")
                cur.execute("ALTER TABLE audit_events ADD COLUMN IF NOT EXISTS prev_hash TEXT")
                cur.execute("CREATE INDEX IF NOT EXISTS idx_audit_session ON audit_events(session_id)")
                cur.execute("CREATE INDEX IF NOT EXISTS idx_audit_timestamp ON audit_events(timestamp)")
                cur.execute("CREATE INDEX IF NOT EXISTS idx_audit_tenant ON audit_events(tenant_id)")
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS audit_chain_anchors (
                        tenant_id     TEXT PRIMARY KEY,
                        anchor_hash   TEXT NOT NULL,
                        purged_before TIMESTAMPTZ NOT NULL,
                        purged_count  INTEGER NOT NULL,
                        created_at    TIMESTAMPTZ NOT NULL
                    )
                """)

    def record(self, event: AuditEvent, tenant_id: str = "default") -> None:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT chain_hash FROM audit_events WHERE tenant_id=%s ORDER BY timestamp DESC LIMIT 1",
                    (tenant_id,),
                )
                row = cur.fetchone()
                prev = (row[0] if row and row[0] else None) or "GENESIS"
                chain_hash = _compute_pg_chain_hash(prev, event, tenant_id)
                cur.execute(
                    """INSERT INTO audit_events
                       (event_id, session_id, agent_role, user_id, source_id,
                        query, decision, policy_name, reason, context_hash,
                        timestamp, tenant_id, sensitivity_level, source_type,
                        catalog_enriched, catalog_connection_id, chain_hash, prev_hash)
                       VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)""",
                    (event.event_id, event.session_id, event.agent_role, event.user_id,
                     event.source_id, event.query, event.decision.value, event.policy_name,
                     event.reason, event.context_hash, event.timestamp, tenant_id,
                     event.sensitivity_level, event.source_type,
                     event.catalog_enriched, event.catalog_connection_id,
                     chain_hash, prev),
                )

    def get_session_events(self, session_id: str, tenant_id: str = "default") -> list[AuditEvent]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    "SELECT * FROM audit_events WHERE session_id=%s AND tenant_id=%s ORDER BY timestamp",
                    (session_id, tenant_id),
                )
                return [self._row_to_event(r) for r in cur.fetchall()]

    def get_all_events(self, tenant_id: str = "default") -> list[AuditEvent]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    "SELECT * FROM audit_events WHERE tenant_id=%s ORDER BY timestamp",
                    (tenant_id,),
                )
                return [self._row_to_event(r) for r in cur.fetchall()]

    def get_session_owner(self, session_id: str, tenant_id: str = "default") -> str | None:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """SELECT agent_role FROM audit_events
                       WHERE session_id=%s AND tenant_id=%s ORDER BY timestamp LIMIT 1""",
                    (session_id, tenant_id),
                )
                row = cur.fetchone()
        return row[0] if row else None

    def _row_to_event(self, row: dict) -> AuditEvent:
        ts = row["timestamp"]
        if isinstance(ts, str):
            ts = datetime.fromisoformat(ts)
        return AuditEvent(
            event_id=row["event_id"], session_id=row["session_id"],
            agent_role=row["agent_role"], user_id=row["user_id"],
            source_id=row["source_id"], query=row["query"],
            decision=Decision(row["decision"]), policy_name=row["policy_name"],
            reason=row["reason"], timestamp=ts, context_hash=row.get("context_hash"),
            sensitivity_level=row.get("sensitivity_level"),
            source_type=row.get("source_type", "rest"),
            catalog_enriched=bool(row.get("catalog_enriched", False)),
            catalog_connection_id=row.get("catalog_connection_id"),
            chain_hash=row.get("chain_hash"),
            prev_hash=row.get("prev_hash"),
        )

    def _get_chain_anchor(self, tenant_id: str) -> dict | None:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    "SELECT anchor_hash, purged_before, purged_count FROM audit_chain_anchors WHERE tenant_id=%s",
                    (tenant_id,),
                )
                row = cur.fetchone()
        if not row:
            return None
        return {"anchor_hash": row["anchor_hash"], "purged_before": str(row["purged_before"]), "purged_count": row["purged_count"]}

    def verify_chain(self, tenant_id: str = "default") -> dict:
        events = self.get_all_events(tenant_id)
        total = len(events)
        chain_start = next((i for i, e in enumerate(events) if e.chain_hash is not None), None)
        if chain_start is None:
            return {"valid": True, "total": total, "chained": 0, "legacy": total, "broken_at": None}
        legacy = chain_start
        chained_events = events[chain_start:]
        # Seed from stored anchor if a purge has occurred, otherwise use GENESIS
        anchor = self._get_chain_anchor(tenant_id)
        prev = anchor["anchor_hash"] if anchor else "GENESIS"
        for idx, event in enumerate(chained_events):
            if event.prev_hash != prev:
                return {
                    "valid": False, "total": total,
                    "chained": idx, "legacy": legacy,
                    "broken_at": event.event_id,
                }
            expected = _compute_pg_chain_hash(prev, event, tenant_id)
            if event.chain_hash != expected:
                return {
                    "valid": False, "total": total,
                    "chained": idx, "legacy": legacy,
                    "broken_at": event.event_id,
                }
            prev = event.chain_hash
        return {"valid": True, "total": total, "chained": len(chained_events), "legacy": legacy, "broken_at": None}

    def purge_old_events(self, retention_days: int, tenant_id: str | None = None) -> int:
        from datetime import timedelta
        cutoff = datetime.utcnow() - timedelta(days=retention_days)
        with self._conn() as conn:
            with conn.cursor() as cur:
                if tenant_id:
                    tenant_ids = [tenant_id]
                else:
                    cur.execute(
                        "SELECT DISTINCT tenant_id FROM audit_events WHERE timestamp < %s", (cutoff,)
                    )
                    tenant_ids = [r[0] for r in cur.fetchall()]

                total_deleted = 0
                for tid in tenant_ids:
                    # Capture chain_hash of the last event being deleted — becomes the anchor
                    cur.execute(
                        """SELECT chain_hash FROM audit_events
                           WHERE tenant_id=%s AND timestamp < %s
                           ORDER BY timestamp DESC LIMIT 1""",
                        (tid, cutoff),
                    )
                    row = cur.fetchone()
                    if not row:
                        continue
                    anchor_hash = row[0] or "GENESIS"

                    cur.execute(
                        "SELECT COUNT(*) FROM audit_events WHERE tenant_id=%s AND timestamp < %s",
                        (tid, cutoff),
                    )
                    purged_count = cur.fetchone()[0]

                    cur.execute(
                        """INSERT INTO audit_chain_anchors
                               (tenant_id, anchor_hash, purged_before, purged_count, created_at)
                           VALUES (%s, %s, %s, %s, %s)
                           ON CONFLICT(tenant_id) DO UPDATE SET
                               anchor_hash=EXCLUDED.anchor_hash,
                               purged_before=EXCLUDED.purged_before,
                               purged_count=EXCLUDED.purged_count,
                               created_at=EXCLUDED.created_at""",
                        (tid, anchor_hash, cutoff, purged_count, datetime.utcnow()),
                    )
                    cur.execute(
                        "DELETE FROM audit_events WHERE tenant_id=%s AND timestamp < %s",
                        (tid, cutoff),
                    )
                    total_deleted += purged_count

        return total_deleted


class PostgresKeyStore(_PostgresBase, KeyStoreBase):

    def __init__(self, database_url: str):
        super().__init__(database_url)
        self._init_db()

    def _init_db(self):
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS api_keys (
                        key_id        TEXT PRIMARY KEY,
                        key_hash      TEXT NOT NULL UNIQUE,
                        name          TEXT NOT NULL,
                        created_at    TIMESTAMPTZ NOT NULL,
                        last_used     TIMESTAMPTZ,
                        is_active     BOOLEAN NOT NULL DEFAULT TRUE,
                        tenant_id     TEXT NOT NULL DEFAULT 'default',
                        is_superadmin BOOLEAN NOT NULL DEFAULT FALSE,
                        scope         TEXT NOT NULL DEFAULT 'admin',
                        expires_at    TIMESTAMPTZ
                    )
                """)
                cur.execute("ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS tenant_id TEXT NOT NULL DEFAULT 'default'")
                cur.execute("ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS is_superadmin BOOLEAN NOT NULL DEFAULT FALSE")
                cur.execute("ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS scope TEXT NOT NULL DEFAULT 'admin'")
                cur.execute("ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS expires_at TIMESTAMPTZ")
                cur.execute("CREATE INDEX IF NOT EXISTS idx_keys_tenant ON api_keys(tenant_id)")

    def create_key(
        self,
        name: str,
        tenant_id: str = "default",
        is_superadmin: bool = False,
        scope: str = "admin",
        expires_days: int | None = None,
    ) -> tuple[str, str]:
        key_id    = secrets.token_hex(8)
        plaintext = KEY_PREFIX + secrets.token_hex(24)
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """INSERT INTO api_keys
                       (key_id, key_hash, name, created_at, tenant_id, is_superadmin, scope, expires_at)
                       VALUES (%s,%s,%s,NOW(),%s,%s,%s,
                               CASE WHEN %s IS NOT NULL
                                    THEN NOW() + (%s || ' days')::INTERVAL
                                    ELSE NULL END)""",
                    (key_id, _hash(plaintext), name, tenant_id, is_superadmin, scope,
                     str(expires_days) if expires_days is not None else None,
                     str(expires_days) if expires_days is not None else None),
                )
        return key_id, plaintext

    def validate(self, plaintext: str) -> dict | None:
        if not plaintext.startswith(KEY_PREFIX):
            return None
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """SELECT key_id, tenant_id, is_superadmin, scope
                       FROM api_keys
                       WHERE key_hash=%s AND is_active=TRUE
                         AND (expires_at IS NULL OR expires_at > NOW())""",
                    (_hash(plaintext),),
                )
                row = cur.fetchone()
                if row:
                    cur.execute("UPDATE api_keys SET last_used=NOW() WHERE key_id=%s", (row[0],))
                    return {
                        "key_id": row[0], "tenant_id": row[1],
                        "is_superadmin": bool(row[2]), "scope": row[3] or "admin",
                    }
        return None

    def list_keys(self, tenant_id: str = "default") -> list[dict]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    """SELECT key_id, name, created_at, last_used, is_active,
                              tenant_id, is_superadmin, scope, expires_at
                       FROM api_keys WHERE tenant_id=%s ORDER BY created_at""",
                    (tenant_id,),
                )
                rows = cur.fetchall()

        def _iso(v):
            return v.isoformat() if v and hasattr(v, "isoformat") else v

        return [
            {
                "key_id": r["key_id"], "name": r["name"],
                "created_at": _iso(r["created_at"]), "last_used": _iso(r["last_used"]),
                "is_active": bool(r["is_active"]), "tenant_id": r["tenant_id"],
                "is_superadmin": bool(r["is_superadmin"]),
                "scope": r["scope"] or "admin", "expires_at": _iso(r["expires_at"]),
            }
            for r in rows
        ]

    def revoke(self, key_id: str, tenant_id: str = "default") -> bool:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "UPDATE api_keys SET is_active=FALSE WHERE key_id=%s AND tenant_id=%s AND is_active=TRUE",
                    (key_id, tenant_id),
                )
                return cur.rowcount > 0

    def rotate_key(self, key_id: str, tenant_id: str) -> tuple[str, str] | None:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    """SELECT name, is_superadmin, scope,
                              CASE WHEN expires_at IS NOT NULL
                                   THEN GREATEST(1, DATE_PART('day', expires_at - NOW())::int)
                                   ELSE NULL END AS expires_days
                       FROM api_keys WHERE key_id=%s AND tenant_id=%s AND is_active=TRUE""",
                    (key_id, tenant_id),
                )
                row = cur.fetchone()
        if not row:
            return None
        new_id, plaintext = self.create_key(
            row["name"], tenant_id, bool(row["is_superadmin"]),
            row["scope"] or "admin", row["expires_days"],
        )
        self.revoke(key_id, tenant_id)
        return new_id, plaintext

    def count(self, tenant_id: str | None = None) -> int:
        with self._conn() as conn:
            with conn.cursor() as cur:
                if tenant_id:
                    cur.execute("SELECT COUNT(*) FROM api_keys WHERE tenant_id=%s", (tenant_id,))
                else:
                    cur.execute("SELECT COUNT(*) FROM api_keys")
                return cur.fetchone()[0]


class PostgresTenantStore(_PostgresBase, TenantStoreBase):

    def __init__(self, database_url: str):
        super().__init__(database_url)
        self._init_db()

    def _init_db(self):
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS tenants (
                        tenant_id  TEXT PRIMARY KEY,
                        name       TEXT NOT NULL UNIQUE,
                        created_at TIMESTAMPTZ NOT NULL,
                        is_active  BOOLEAN NOT NULL DEFAULT TRUE
                    )
                """)

    def create_tenant(self, name: str) -> str:
        tenant_id = "ten_" + secrets.token_hex(8)
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "INSERT INTO tenants (tenant_id, name, created_at) VALUES (%s,%s,NOW())",
                    (tenant_id, name),
                )
        return tenant_id

    def get_tenant(self, tenant_id: str) -> dict | None:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    "SELECT tenant_id, name, created_at, is_active FROM tenants WHERE tenant_id=%s",
                    (tenant_id,),
                )
                row = cur.fetchone()
        if not row:
            return None
        return {"tenant_id": row["tenant_id"], "name": row["name"],
                "created_at": row["created_at"].isoformat() if hasattr(row["created_at"], "isoformat") else row["created_at"],
                "is_active": bool(row["is_active"])}

    def list_tenants(self) -> list[dict]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute("SELECT tenant_id, name, created_at, is_active FROM tenants ORDER BY created_at")
                rows = cur.fetchall()
        return [{"tenant_id": r["tenant_id"], "name": r["name"],
                 "created_at": r["created_at"].isoformat() if hasattr(r["created_at"], "isoformat") else r["created_at"],
                 "is_active": bool(r["is_active"])} for r in rows]

    def deactivate_tenant(self, tenant_id: str) -> bool:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "UPDATE tenants SET is_active=FALSE WHERE tenant_id=%s AND is_active=TRUE",
                    (tenant_id,),
                )
                return cur.rowcount > 0

    def count(self) -> int:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("SELECT COUNT(*) FROM tenants")
                return cur.fetchone()[0]


import json


class PostgresPolicyStore(_PostgresBase, PolicyStoreBase):

    def __init__(self, database_url: str):
        super().__init__(database_url)
        self._init_db()

    def _init_db(self):
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS policies (
                        policy_id       TEXT PRIMARY KEY,
                        name            TEXT NOT NULL,
                        agent_role      TEXT NOT NULL,
                        allowed_sources TEXT NOT NULL DEFAULT '[]',
                        denied_sources  TEXT NOT NULL DEFAULT '[]',
                        allowed_tasks   TEXT NOT NULL DEFAULT '[]',
                        denied_tasks    TEXT NOT NULL DEFAULT '[]',
                        max_sensitivity TEXT NOT NULL DEFAULT 'high',
                        industry        TEXT,
                        category        TEXT,
                        tenant_id       TEXT NOT NULL,
                        created_at      TIMESTAMPTZ NOT NULL,
                        updated_at      TIMESTAMPTZ NOT NULL,
                        version         INTEGER NOT NULL DEFAULT 1,
                        changed_by      TEXT,
                        is_active       BOOLEAN NOT NULL DEFAULT TRUE,
                        UNIQUE(name, tenant_id)
                    )
                """)
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS policy_versions (
                        version_id      TEXT PRIMARY KEY,
                        policy_id       TEXT NOT NULL,
                        name            TEXT NOT NULL,
                        agent_role      TEXT NOT NULL,
                        allowed_sources TEXT NOT NULL DEFAULT '[]',
                        denied_sources  TEXT NOT NULL DEFAULT '[]',
                        allowed_tasks   TEXT NOT NULL DEFAULT '[]',
                        denied_tasks    TEXT NOT NULL DEFAULT '[]',
                        max_sensitivity TEXT NOT NULL,
                        industry        TEXT,
                        category        TEXT,
                        version         INTEGER NOT NULL,
                        tenant_id       TEXT NOT NULL,
                        changed_by      TEXT,
                        changed_at      TIMESTAMPTZ NOT NULL
                    )
                """)
                # Migrations: add columns if upgrading from older schema
                for col, definition in [
                    ("allowed_tasks", "TEXT NOT NULL DEFAULT '[]'"),
                    ("denied_tasks",  "TEXT NOT NULL DEFAULT '[]'"),
                    ("industry",      "TEXT"),
                    ("category",      "TEXT"),
                ]:
                    cur.execute(f"ALTER TABLE policies ADD COLUMN IF NOT EXISTS {col} {definition}")
                    cur.execute(f"ALTER TABLE policy_versions ADD COLUMN IF NOT EXISTS {col} {definition}")

    def upsert(self, policy: dict, tenant_id: str, changed_by: str | None = None) -> str:
        name            = policy["name"]
        agent_role      = policy["agent_role"]
        allowed         = json.dumps(policy.get("allowed_sources", []))
        denied          = json.dumps(policy.get("denied_sources", []))
        allowed_tasks   = json.dumps(policy.get("allowed_tasks", []))
        denied_tasks    = json.dumps(policy.get("denied_tasks", []))
        sensitivity     = policy.get("max_sensitivity", "high")
        industry        = policy.get("industry")
        category        = policy.get("category")

        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT policy_id, version FROM policies WHERE name=%s AND tenant_id=%s AND is_active=TRUE",
                    (name, tenant_id),
                )
                existing = cur.fetchone()

                if existing:
                    policy_id, version = existing
                    new_version = version + 1
                    cur.execute(
                        """UPDATE policies SET agent_role=%s, allowed_sources=%s, denied_sources=%s,
                           allowed_tasks=%s, denied_tasks=%s, max_sensitivity=%s,
                           industry=%s, category=%s, updated_at=NOW(), version=%s, changed_by=%s
                           WHERE policy_id=%s""",
                        (agent_role, allowed, denied, allowed_tasks, denied_tasks,
                         sensitivity, industry, category, new_version, changed_by, policy_id),
                    )
                    cur.execute(
                        """INSERT INTO policy_versions
                           (version_id, policy_id, name, agent_role, allowed_sources, denied_sources,
                            allowed_tasks, denied_tasks, max_sensitivity, industry, category,
                            version, tenant_id, changed_by, changed_at)
                           VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,NOW())""",
                        (secrets.token_hex(8), policy_id, name, agent_role, allowed, denied,
                         allowed_tasks, denied_tasks, sensitivity, industry, category,
                         new_version, tenant_id, changed_by),
                    )
                else:
                    policy_id = "pol_" + secrets.token_hex(8)
                    cur.execute(
                        """INSERT INTO policies
                           (policy_id, name, agent_role, allowed_sources, denied_sources,
                            allowed_tasks, denied_tasks, max_sensitivity, industry, category,
                            tenant_id, created_at, updated_at, version, changed_by)
                           VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,NOW(),NOW(),1,%s)""",
                        (policy_id, name, agent_role, allowed, denied, allowed_tasks, denied_tasks,
                         sensitivity, industry, category, tenant_id, changed_by),
                    )
                    cur.execute(
                        """INSERT INTO policy_versions
                           (version_id, policy_id, name, agent_role, allowed_sources, denied_sources,
                            allowed_tasks, denied_tasks, max_sensitivity, industry, category,
                            version, tenant_id, changed_by, changed_at)
                           VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,1,%s,%s,NOW())""",
                        (secrets.token_hex(8), policy_id, name, agent_role, allowed, denied,
                         allowed_tasks, denied_tasks, sensitivity, industry, category, tenant_id, changed_by),
                    )
        return policy_id

    def get(self, name: str, tenant_id: str) -> dict | None:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    """SELECT policy_id, name, agent_role, allowed_sources, denied_sources,
                              allowed_tasks, denied_tasks, max_sensitivity, industry, category,
                              version, created_at, updated_at, changed_by
                       FROM policies WHERE name=%s AND tenant_id=%s AND is_active=TRUE""",
                    (name, tenant_id),
                )
                row = cur.fetchone()
        return self._row_to_dict(row) if row else None

    def list_active(self, tenant_id: str) -> list[dict]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    """SELECT policy_id, name, agent_role, allowed_sources, denied_sources,
                              allowed_tasks, denied_tasks, max_sensitivity, industry, category,
                              version, created_at, updated_at, changed_by
                       FROM policies WHERE tenant_id=%s AND is_active=TRUE ORDER BY industry, category, name""",
                    (tenant_id,),
                )
                return [self._row_to_dict(r) for r in cur.fetchall()]

    def delete(self, name: str, tenant_id: str) -> bool:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "UPDATE policies SET is_active=FALSE WHERE name=%s AND tenant_id=%s AND is_active=TRUE",
                    (name, tenant_id),
                )
                return cur.rowcount > 0

    def get_history(self, name: str, tenant_id: str) -> list[dict]:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT policy_id FROM policies WHERE name=%s AND tenant_id=%s",
                    (name, tenant_id),
                )
                row = cur.fetchone()
            if not row:
                return []
            policy_id = row[0]
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    """SELECT version_id, version, agent_role, allowed_sources, denied_sources,
                              allowed_tasks, denied_tasks, max_sensitivity, industry, category,
                              changed_by, changed_at
                       FROM policy_versions WHERE policy_id=%s AND tenant_id=%s
                       ORDER BY version DESC""",
                    (policy_id, tenant_id),
                )
                return [
                    {
                        "version_id":      r["version_id"],
                        "version":         r["version"],
                        "agent_role":      r["agent_role"],
                        "allowed_sources": json.loads(r["allowed_sources"]),
                        "denied_sources":  json.loads(r["denied_sources"]),
                        "allowed_tasks":   json.loads(r["allowed_tasks"]) if r["allowed_tasks"] else [],
                        "denied_tasks":    json.loads(r["denied_tasks"]) if r["denied_tasks"] else [],
                        "max_sensitivity": r["max_sensitivity"],
                        "industry":        r["industry"],
                        "category":        r["category"],
                        "changed_by":      r["changed_by"],
                        "changed_at":      r["changed_at"].isoformat() if hasattr(r["changed_at"], "isoformat") else r["changed_at"],
                    }
                    for r in cur.fetchall()
                ]

    def count(self, tenant_id: str) -> int:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT COUNT(*) FROM policies WHERE tenant_id=%s AND is_active=TRUE",
                    (tenant_id,),
                )
                return cur.fetchone()[0]

    def _row_to_dict(self, row) -> dict:
        def iso(v):
            return v.isoformat() if hasattr(v, "isoformat") else v
        return {
            "policy_id":      row["policy_id"],
            "name":           row["name"],
            "agent_role":     row["agent_role"],
            "allowed_sources": json.loads(row["allowed_sources"]),
            "denied_sources":  json.loads(row["denied_sources"]),
            "allowed_tasks":   json.loads(row["allowed_tasks"]) if row.get("allowed_tasks") else [],
            "denied_tasks":    json.loads(row["denied_tasks"]) if row.get("denied_tasks") else [],
            "max_sensitivity": row["max_sensitivity"],
            "industry":        row.get("industry"),
            "category":        row.get("category"),
            "version":         row["version"],
            "created_at":      iso(row["created_at"]),
            "updated_at":      iso(row["updated_at"]),
            "changed_by":      row["changed_by"],
        }


class PostgresAlertStore(_PostgresBase, AlertStoreBase):

    def __init__(self, database_url: str):
        super().__init__(database_url)
        self._init_db()

    def _init_db(self):
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS alert_rules (
                        rule_id          TEXT PRIMARY KEY,
                        tenant_id        TEXT NOT NULL,
                        name             TEXT NOT NULL,
                        rule_type        TEXT NOT NULL,
                        threshold        INTEGER NOT NULL DEFAULT 5,
                        window_minutes   INTEGER NOT NULL DEFAULT 60,
                        cooldown_minutes INTEGER NOT NULL DEFAULT 60,
                        webhook_url      TEXT,
                        email            TEXT,
                        is_enabled       BOOLEAN NOT NULL DEFAULT TRUE,
                        created_at       TIMESTAMPTZ NOT NULL,
                        last_fired_at    TIMESTAMPTZ
                    )
                """)
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS alert_deliveries (
                        delivery_id TEXT PRIMARY KEY,
                        rule_id     TEXT NOT NULL,
                        tenant_id   TEXT NOT NULL,
                        fired_at    TIMESTAMPTZ NOT NULL,
                        rule_name   TEXT NOT NULL,
                        rule_type   TEXT NOT NULL,
                        detail      TEXT NOT NULL,
                        status      TEXT NOT NULL,
                        error       TEXT
                    )
                """)
                cur.execute("CREATE INDEX IF NOT EXISTS idx_alert_rules_tenant ON alert_rules(tenant_id)")
                cur.execute("CREATE INDEX IF NOT EXISTS idx_alert_deliveries_tenant ON alert_deliveries(tenant_id)")

    def create_rule(self, rule: dict, tenant_id: str) -> str:
        rule_id = "alr_" + secrets.token_hex(8)
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """INSERT INTO alert_rules
                       (rule_id, tenant_id, name, rule_type, threshold, window_minutes,
                        cooldown_minutes, webhook_url, email, is_enabled, created_at)
                       VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,TRUE,NOW())""",
                    (rule_id, tenant_id, rule["name"], rule["rule_type"],
                     rule.get("threshold", 5), rule.get("window_minutes", 60),
                     rule.get("cooldown_minutes", 60),
                     rule.get("webhook_url"), rule.get("email")),
                )
        return rule_id

    def get_rule(self, rule_id: str, tenant_id: str) -> dict | None:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    "SELECT * FROM alert_rules WHERE rule_id=%s AND tenant_id=%s",
                    (rule_id, tenant_id),
                )
                row = cur.fetchone()
        return self._rule_row(dict(row)) if row else None

    def list_rules(self, tenant_id: str) -> list[dict]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    "SELECT * FROM alert_rules WHERE tenant_id=%s ORDER BY created_at",
                    (tenant_id,),
                )
                return [self._rule_row(dict(r)) for r in cur.fetchall()]

    def update_rule(self, rule_id: str, updates: dict, tenant_id: str) -> bool:
        allowed = {"name", "threshold", "window_minutes", "cooldown_minutes",
                   "webhook_url", "email", "is_enabled"}
        fields = {k: v for k, v in updates.items() if k in allowed}
        if not fields:
            return False
        set_clause = ", ".join(f"{k}=%s" for k in fields)
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    f"UPDATE alert_rules SET {set_clause} WHERE rule_id=%s AND tenant_id=%s",
                    (*fields.values(), rule_id, tenant_id),
                )
                return cur.rowcount > 0

    def delete_rule(self, rule_id: str, tenant_id: str) -> bool:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "DELETE FROM alert_rules WHERE rule_id=%s AND tenant_id=%s",
                    (rule_id, tenant_id),
                )
                return cur.rowcount > 0

    def mark_fired(self, rule_id: str, fired_at: str) -> None:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "UPDATE alert_rules SET last_fired_at=%s WHERE rule_id=%s",
                    (fired_at, rule_id),
                )

    def record_delivery(self, delivery: dict) -> None:
        delivery_id = secrets.token_hex(12)
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """INSERT INTO alert_deliveries
                       (delivery_id, rule_id, tenant_id, fired_at, rule_name,
                        rule_type, detail, status, error)
                       VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s)""",
                    (delivery_id, delivery["rule_id"], delivery["tenant_id"],
                     delivery["fired_at"], delivery["rule_name"], delivery["rule_type"],
                     delivery["detail"], delivery["status"], delivery.get("error")),
                )

    def list_deliveries(self, tenant_id: str, limit: int = 50) -> list[dict]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    """SELECT delivery_id, rule_id, tenant_id, fired_at, rule_name,
                              rule_type, detail, status, error
                       FROM alert_deliveries WHERE tenant_id=%s
                       ORDER BY fired_at DESC LIMIT %s""",
                    (tenant_id, limit),
                )
                rows = cur.fetchall()
        return [
            {
                "delivery_id": r["delivery_id"], "rule_id": r["rule_id"],
                "tenant_id": r["tenant_id"],
                "fired_at": r["fired_at"].isoformat() if hasattr(r["fired_at"], "isoformat") else r["fired_at"],
                "rule_name": r["rule_name"], "rule_type": r["rule_type"],
                "detail": json.loads(r["detail"]), "status": r["status"], "error": r["error"],
            }
            for r in rows
        ]

    def _rule_row(self, row: dict) -> dict:
        def iso(v):
            return v.isoformat() if hasattr(v, "isoformat") else v
        return {
            "rule_id": row["rule_id"], "tenant_id": row["tenant_id"],
            "name": row["name"], "rule_type": row["rule_type"],
            "threshold": row["threshold"], "window_minutes": row["window_minutes"],
            "cooldown_minutes": row["cooldown_minutes"],
            "webhook_url": row.get("webhook_url"), "email": row.get("email"),
            "is_enabled": bool(row["is_enabled"]),
            "created_at": iso(row["created_at"]),
            "last_fired_at": iso(row["last_fired_at"]) if row.get("last_fired_at") else None,
        }


class PostgresCatalogStore(_PostgresBase, CatalogStoreBase):
    """
    Postgres-backed catalog store for production deployments.

    Three tables mirror the SQLite schema exactly:
      catalog_connections       — registered catalog integrations (credentials encrypted at rest)
      catalog_classifications   — cached sensitivity classifications (unique per connection + asset_fqn)
      catalog_sync_history      — sync run records

    All timestamps are TIMESTAMPTZ; upserts use ON CONFLICT for idempotency.
    """

    def __init__(self, database_url: str):
        super().__init__(database_url)
        self._init_db()

    def _init_db(self) -> None:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS catalog_connections (
                        connection_id         TEXT PRIMARY KEY,
                        tenant_id             TEXT NOT NULL,
                        catalog_type          TEXT NOT NULL,
                        display_name          TEXT NOT NULL,
                        base_url              TEXT NOT NULL,
                        credentials_enc       TEXT NOT NULL,
                        sync_interval_minutes INTEGER NOT NULL DEFAULT 15,
                        is_enabled            BOOLEAN NOT NULL DEFAULT TRUE,
                        created_at            TIMESTAMPTZ NOT NULL,
                        updated_at            TIMESTAMPTZ NOT NULL,
                        last_sync_at          TIMESTAMPTZ,
                        webhook_secret        TEXT,
                        status                TEXT NOT NULL DEFAULT 'active'
                    )
                """)
                cur.execute(
                    "CREATE INDEX IF NOT EXISTS idx_catalog_conn_tenant ON catalog_connections(tenant_id)"
                )
                # Migration: add status column to existing databases
                cur.execute("""
                    DO $$
                    BEGIN
                        IF NOT EXISTS (
                            SELECT 1 FROM information_schema.columns
                            WHERE table_name='catalog_connections' AND column_name='status'
                        ) THEN
                            ALTER TABLE catalog_connections
                            ADD COLUMN status TEXT NOT NULL DEFAULT 'active';
                        END IF;
                    END$$;
                """)

                cur.execute("""
                    CREATE TABLE IF NOT EXISTS catalog_classifications (
                        classification_id  TEXT NOT NULL,
                        connection_id      TEXT NOT NULL,
                        tenant_id          TEXT NOT NULL,
                        asset_fqn          TEXT NOT NULL,
                        asset_type         TEXT NOT NULL DEFAULT 'table',
                        sensitivity_label  TEXT NOT NULL,
                        raw_label          TEXT NOT NULL,
                        data_classes       TEXT NOT NULL DEFAULT '[]',
                        is_certified       BOOLEAN,
                        trust_score        REAL,
                        catalog_policy_ids TEXT NOT NULL DEFAULT '[]',
                        fetched_at         TIMESTAMPTZ NOT NULL,
                        expires_at         TIMESTAMPTZ NOT NULL,
                        PRIMARY KEY (connection_id, asset_fqn)
                    )
                """)
                cur.execute(
                    "CREATE INDEX IF NOT EXISTS idx_cls_pg_tenant_fqn ON catalog_classifications(tenant_id, asset_fqn)"
                )
                cur.execute(
                    "CREATE INDEX IF NOT EXISTS idx_cls_pg_conn ON catalog_classifications(connection_id)"
                )
                cur.execute(
                    "CREATE INDEX IF NOT EXISTS idx_cls_pg_expires ON catalog_classifications(expires_at)"
                )

                cur.execute("""
                    CREATE TABLE IF NOT EXISTS catalog_sync_history (
                        sync_id        TEXT PRIMARY KEY,
                        connection_id  TEXT NOT NULL,
                        tenant_id      TEXT NOT NULL,
                        status         TEXT NOT NULL,
                        assets_synced  INTEGER NOT NULL DEFAULT 0,
                        assets_skipped INTEGER NOT NULL DEFAULT 0,
                        error_message  TEXT,
                        started_at     TIMESTAMPTZ NOT NULL,
                        completed_at   TIMESTAMPTZ
                    )
                """)
                cur.execute(
                    "CREATE INDEX IF NOT EXISTS idx_sync_pg_conn ON catalog_sync_history(connection_id, started_at)"
                )

    # ── helpers ───────────────────────────────────────────────────────────────

    @staticmethod
    def _iso(v) -> str | None:
        if v is None:
            return None
        return v.isoformat() if hasattr(v, "isoformat") else v

    def _row_to_conn(self, row: dict) -> CatalogConnection:
        return CatalogConnection(
            connection_id=row["connection_id"],
            tenant_id=row["tenant_id"],
            catalog_type=row["catalog_type"],
            display_name=row["display_name"],
            base_url=row["base_url"],
            credentials_enc=row["credentials_enc"],
            sync_interval_minutes=row["sync_interval_minutes"],
            is_enabled=bool(row["is_enabled"]),
            created_at=row["created_at"] if not isinstance(row["created_at"], str)
                       else datetime.fromisoformat(row["created_at"]),
            updated_at=row["updated_at"] if not isinstance(row["updated_at"], str)
                       else datetime.fromisoformat(row["updated_at"]),
            last_sync_at=row.get("last_sync_at"),
            webhook_secret=row.get("webhook_secret"),
            status=row.get("status") or "active",
        )

    def _row_to_cls(self, row: dict) -> CatalogClassification:
        return CatalogClassification(
            classification_id=row["classification_id"],
            connection_id=row["connection_id"],
            tenant_id=row["tenant_id"],
            asset_fqn=row["asset_fqn"],
            asset_type=row["asset_type"],
            sensitivity_label=row["sensitivity_label"],
            raw_label=row["raw_label"],
            data_classes=json.loads(row["data_classes"]),
            is_certified=bool(row["is_certified"]) if row.get("is_certified") is not None else None,
            trust_score=row.get("trust_score"),
            catalog_policy_ids=json.loads(row["catalog_policy_ids"]),
            fetched_at=row["fetched_at"] if not isinstance(row["fetched_at"], str)
                       else datetime.fromisoformat(row["fetched_at"]),
            expires_at=row["expires_at"] if not isinstance(row["expires_at"], str)
                       else datetime.fromisoformat(row["expires_at"]),
        )

    def _row_to_sync(self, row: dict) -> SyncResult:
        return SyncResult(
            sync_id=row["sync_id"],
            connection_id=row["connection_id"],
            tenant_id=row["tenant_id"],
            status=row["status"],
            assets_synced=row["assets_synced"],
            assets_skipped=row["assets_skipped"],
            error_message=row.get("error_message"),
            started_at=row["started_at"] if not isinstance(row["started_at"], str)
                       else datetime.fromisoformat(row["started_at"]),
            completed_at=row.get("completed_at"),
        )

    # ── connections ───────────────────────────────────────────────────────────

    def create_connection(self, conn_obj: CatalogConnection) -> str:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """INSERT INTO catalog_connections
                       (connection_id, tenant_id, catalog_type, display_name, base_url,
                        credentials_enc, sync_interval_minutes, is_enabled, created_at,
                        updated_at, last_sync_at, webhook_secret, status)
                       VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)""",
                    (
                        conn_obj.connection_id, conn_obj.tenant_id, conn_obj.catalog_type,
                        conn_obj.display_name, conn_obj.base_url, conn_obj.credentials_enc,
                        conn_obj.sync_interval_minutes, conn_obj.is_enabled,
                        conn_obj.created_at, conn_obj.updated_at,
                        conn_obj.last_sync_at, conn_obj.webhook_secret,
                        conn_obj.status,
                    ),
                )
        return conn_obj.connection_id

    def get_connection(self, connection_id: str, tenant_id: str) -> CatalogConnection | None:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    "SELECT * FROM catalog_connections WHERE connection_id=%s AND tenant_id=%s",
                    (connection_id, tenant_id),
                )
                row = cur.fetchone()
        return self._row_to_conn(dict(row)) if row else None

    def list_connections(self, tenant_id: str) -> list[CatalogConnection]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    "SELECT * FROM catalog_connections WHERE tenant_id=%s ORDER BY created_at",
                    (tenant_id,),
                )
                return [self._row_to_conn(dict(r)) for r in cur.fetchall()]

    def list_all_connections(self) -> list[CatalogConnection]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    "SELECT * FROM catalog_connections ORDER BY tenant_id, created_at"
                )
                return [self._row_to_conn(dict(r)) for r in cur.fetchall()]

    def update_connection(self, connection_id: str, updates: dict, tenant_id: str) -> bool:
        allowed = {"display_name", "sync_interval_minutes", "is_enabled", "webhook_secret", "credentials_enc", "status"}
        fields = {k: v for k, v in updates.items() if k in allowed}
        if not fields:
            return False
        fields["updated_at"] = datetime.utcnow()
        set_clause = ", ".join(f"{k}=%s" for k in fields)
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    f"UPDATE catalog_connections SET {set_clause} WHERE connection_id=%s AND tenant_id=%s",
                    (*fields.values(), connection_id, tenant_id),
                )
                return cur.rowcount > 0

    def delete_connection(self, connection_id: str, tenant_id: str) -> bool:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "DELETE FROM catalog_connections WHERE connection_id=%s AND tenant_id=%s",
                    (connection_id, tenant_id),
                )
                return cur.rowcount > 0

    def mark_sync_completed(self, connection_id: str, sync_at: datetime) -> None:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "UPDATE catalog_connections SET last_sync_at=%s WHERE connection_id=%s",
                    (sync_at, connection_id),
                )

    # ── classifications ───────────────────────────────────────────────────────

    def upsert_classifications(self, classifications: list[CatalogClassification]) -> int:
        written = 0
        with self._conn() as conn:
            with conn.cursor() as cur:
                for cls in classifications:
                    cur.execute(
                        """INSERT INTO catalog_classifications
                           (classification_id, connection_id, tenant_id, asset_fqn, asset_type,
                            sensitivity_label, raw_label, data_classes, is_certified, trust_score,
                            catalog_policy_ids, fetched_at, expires_at)
                           VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                           ON CONFLICT (connection_id, asset_fqn) DO UPDATE SET
                               classification_id  = EXCLUDED.classification_id,
                               asset_type         = EXCLUDED.asset_type,
                               sensitivity_label  = EXCLUDED.sensitivity_label,
                               raw_label          = EXCLUDED.raw_label,
                               data_classes       = EXCLUDED.data_classes,
                               is_certified       = EXCLUDED.is_certified,
                               trust_score        = EXCLUDED.trust_score,
                               catalog_policy_ids = EXCLUDED.catalog_policy_ids,
                               fetched_at         = EXCLUDED.fetched_at,
                               expires_at         = EXCLUDED.expires_at""",
                        (
                            cls.classification_id, cls.connection_id, cls.tenant_id,
                            cls.asset_fqn, cls.asset_type, cls.sensitivity_label, cls.raw_label,
                            json.dumps(cls.data_classes),
                            cls.is_certified,
                            cls.trust_score,
                            json.dumps(cls.catalog_policy_ids),
                            cls.fetched_at, cls.expires_at,
                        ),
                    )
                    written += 1
        return written

    def get_classifications_for_asset(
        self, tenant_id: str, asset_fqn: str
    ) -> list[CatalogClassification]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    """SELECT * FROM catalog_classifications
                       WHERE tenant_id=%s AND expires_at > NOW()
                       AND (asset_fqn = %s OR asset_fqn LIKE %s OR %s LIKE asset_fqn || '%%')
                       ORDER BY fetched_at DESC""",
                    (tenant_id, asset_fqn, f"%{asset_fqn}%", asset_fqn),
                )
                return [self._row_to_cls(dict(r)) for r in cur.fetchall()]

    def list_classifications(
        self,
        tenant_id: str,
        connection_id: str | None = None,
        sensitivity_label: str | None = None,
    ) -> list[CatalogClassification]:
        query = "SELECT * FROM catalog_classifications WHERE tenant_id=%s"
        params: list = [tenant_id]
        if connection_id:
            query += " AND connection_id=%s"
            params.append(connection_id)
        if sensitivity_label:
            query += " AND sensitivity_label=%s"
            params.append(sensitivity_label)
        query += " ORDER BY asset_fqn"
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(query, params)
                return [self._row_to_cls(dict(r)) for r in cur.fetchall()]

    def delete_stale_classifications(self, connection_id: str) -> int:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "DELETE FROM catalog_classifications WHERE connection_id=%s AND expires_at <= NOW()",
                    (connection_id,),
                )
                return cur.rowcount

    # ── sync history ──────────────────────────────────────────────────────────

    def record_sync(self, result: SyncResult) -> None:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """INSERT INTO catalog_sync_history
                       (sync_id, connection_id, tenant_id, status, assets_synced,
                        assets_skipped, error_message, started_at, completed_at)
                       VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s)""",
                    (
                        result.sync_id, result.connection_id, result.tenant_id,
                        result.status, result.assets_synced, result.assets_skipped,
                        result.error_message, result.started_at, result.completed_at,
                    ),
                )

    def list_sync_history(
        self, connection_id: str, tenant_id: str, limit: int = 25
    ) -> list[SyncResult]:
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    """SELECT * FROM catalog_sync_history
                       WHERE connection_id=%s AND tenant_id=%s
                       ORDER BY started_at DESC LIMIT %s""",
                    (connection_id, tenant_id, limit),
                )
                return [self._row_to_sync(dict(r)) for r in cur.fetchall()]

    def get_catalog_stats(self, tenant_id: str) -> dict:
        with self._conn() as conn:
            with conn.cursor() as cur:
                # Connection counts
                cur.execute(
                    "SELECT status, is_enabled, COUNT(*) FROM catalog_connections WHERE tenant_id=%s GROUP BY status, is_enabled",
                    (tenant_id,),
                )
                conn_rows = cur.fetchall()
                total_conn    = sum(r[2] for r in conn_rows)
                active_conn   = sum(r[2] for r in conn_rows if r[0] == "active"   and r[1])
                degraded_conn = sum(r[2] for r in conn_rows if r[0] == "degraded")
                disabled_conn = sum(r[2] for r in conn_rows if not r[1])

                # Total classifications
                cur.execute(
                    "SELECT COUNT(*) FROM catalog_classifications WHERE tenant_id=%s",
                    (tenant_id,),
                )
                total_cls = cur.fetchone()[0]

                # Sensitivity distribution
                cur.execute(
                    "SELECT sensitivity_label, COUNT(*) FROM catalog_classifications WHERE tenant_id=%s GROUP BY sensitivity_label",
                    (tenant_id,),
                )
                sensitivity_dist = {r[0]: r[1] for r in cur.fetchall()}

                # Coverage by catalog type
                cur.execute(
                    """SELECT cc.catalog_type, COUNT(cls.classification_id)
                       FROM catalog_connections cc
                       LEFT JOIN catalog_classifications cls ON cc.connection_id = cls.connection_id
                       WHERE cc.tenant_id=%s
                       GROUP BY cc.catalog_type""",
                    (tenant_id,),
                )
                coverage_by_type = {r[0]: r[1] for r in cur.fetchall()}

        return {
            "total_connections":        total_conn,
            "active_connections":       active_conn,
            "degraded_connections":     degraded_conn,
            "disabled_connections":     disabled_conn,
            "total_classifications":    total_cls,
            "sensitivity_distribution": sensitivity_dist,
            "coverage_by_catalog_type": coverage_by_type,
        }


class PostgresTrialStore(_PostgresBase, TrialStoreBase):

    def __init__(self, database_url: str):
        super().__init__(database_url)
        self._init_db()

    def _init_db(self):
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS trials (
                        trial_id          TEXT PRIMARY KEY,
                        tenant_id         TEXT NOT NULL UNIQUE,
                        name              TEXT NOT NULL,
                        email             TEXT NOT NULL,
                        company           TEXT NOT NULL,
                        industry          TEXT NOT NULL DEFAULT 'financial_services',
                        trial_start       TIMESTAMPTZ NOT NULL,
                        trial_end         TIMESTAMPTZ NOT NULL,
                        status            TEXT NOT NULL DEFAULT 'active',
                        email_day7_sent   BOOLEAN NOT NULL DEFAULT FALSE,
                        email_day13_sent  BOOLEAN NOT NULL DEFAULT FALSE
                    )
                """)

    def create_trial(self, trial_id, tenant_id, name, email, company, industry, trial_end):
        from datetime import datetime, timezone
        now = datetime.now(timezone.utc).isoformat()
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "INSERT INTO trials (trial_id,tenant_id,name,email,company,industry,trial_start,trial_end) VALUES (%s,%s,%s,%s,%s,%s,%s,%s)",
                    (trial_id, tenant_id, name, email, company, industry, now, trial_end),
                )

    def get_trial_by_tenant(self, tenant_id):
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute("SELECT * FROM trials WHERE tenant_id = %s", (tenant_id,))
                row = cur.fetchone()
                return dict(row) if row else None

    def list_active_trials(self):
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute("SELECT * FROM trials WHERE status='active'")
                return [dict(r) for r in cur.fetchall()]

    def list_all_trials(self):
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute("SELECT * FROM trials ORDER BY trial_start DESC")
                return [dict(r) for r in cur.fetchall()]

    def expire_trials(self, before):
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    "UPDATE trials SET status='expired' WHERE status='active' AND trial_end < %s RETURNING tenant_id",
                    (before.isoformat(),),
                )
                return [r["tenant_id"] for r in cur.fetchall()]

    def mark_email_sent(self, trial_id, email_type):
        col = "email_day7_sent" if email_type == "day7" else "email_day13_sent"
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(f"UPDATE trials SET {col}=TRUE WHERE trial_id=%s", (trial_id,))
