"""
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, Decision

from .base import AuditLogBase, KeyStoreBase, TenantStoreBase, PolicyStoreBase

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'
                    )
                """)
                # Migration: add tenant_id if upgrading
                cur.execute("""
                    ALTER TABLE audit_events
                    ADD COLUMN IF NOT EXISTS tenant_id TEXT NOT NULL DEFAULT 'default'
                """)
                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)")

    def record(self, event: AuditEvent, tenant_id: str = "default") -> None:
        with self._conn() as conn:
            with conn.cursor() as cur:
                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)
                       VALUES (%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),
                )

    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["context_hash"],
        )


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
                    )
                """)
                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("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) -> 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) VALUES (%s,%s,%s,NOW(),%s,%s)",
                    (key_id, _hash(plaintext), name, tenant_id, is_superadmin),
                )
        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 FROM api_keys WHERE key_hash=%s AND is_active=TRUE",
                    (_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])}
        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 FROM api_keys WHERE tenant_id=%s ORDER BY created_at",
                    (tenant_id,),
                )
                rows = cur.fetchall()
        return [
            {"key_id": r["key_id"], "name": r["name"],
             "created_at": r["created_at"].isoformat() if hasattr(r["created_at"], "isoformat") else r["created_at"],
             "last_used": r["last_used"].isoformat() if r["last_used"] and hasattr(r["last_used"], "isoformat") else r["last_used"],
             "is_active": bool(r["is_active"]), "tenant_id": r["tenant_id"], "is_superadmin": bool(r["is_superadmin"])}
            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 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 '[]',
                        max_sensitivity TEXT NOT NULL DEFAULT 'high',
                        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 '[]',
                        max_sensitivity TEXT NOT NULL,
                        version         INTEGER NOT NULL,
                        tenant_id       TEXT NOT NULL,
                        changed_by      TEXT,
                        changed_at      TIMESTAMPTZ NOT NULL
                    )
                """)

    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", []))
        sensitivity = policy.get("max_sensitivity", "high")

        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,
                           max_sensitivity=%s, updated_at=NOW(), version=%s, changed_by=%s
                           WHERE policy_id=%s""",
                        (agent_role, allowed, denied, sensitivity, new_version, changed_by, policy_id),
                    )
                    # record the new version in history
                    cur.execute(
                        """INSERT INTO policy_versions
                           (version_id, policy_id, name, agent_role, allowed_sources,
                            denied_sources, max_sensitivity, version, tenant_id, changed_by, changed_at)
                           VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,NOW())""",
                        (secrets.token_hex(8), policy_id, name, agent_role, allowed, denied,
                         sensitivity, 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,
                            max_sensitivity, tenant_id, created_at, updated_at, version, changed_by)
                           VALUES (%s,%s,%s,%s,%s,%s,%s,NOW(),NOW(),1,%s)""",
                        (policy_id, name, agent_role, allowed, denied, sensitivity, tenant_id, changed_by),
                    )
                    cur.execute(
                        """INSERT INTO policy_versions
                           (version_id, policy_id, name, agent_role, allowed_sources,
                            denied_sources, max_sensitivity, version, tenant_id, changed_by, changed_at)
                           VALUES (%s,%s,%s,%s,%s,%s,%s,1,%s,%s,NOW())""",
                        (secrets.token_hex(8), policy_id, name, agent_role, allowed, denied, sensitivity, 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,
                              max_sensitivity, 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,
                              max_sensitivity, version, created_at, updated_at, changed_by
                       FROM policies WHERE tenant_id=%s AND is_active=TRUE ORDER BY 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,
                              max_sensitivity, 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"]),
                        "max_sensitivity": r["max_sensitivity"],
                        "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"]),
            "max_sensitivity": row["max_sensitivity"],
            "version": row["version"],
            "created_at": iso(row["created_at"]),
            "updated_at": iso(row["updated_at"]),
            "changed_by": row["changed_by"],
        }
