"""PostgreSQL-backed event bus implementation."""

from __future__ import annotations

import asyncio
import json
import logging
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Iterable
from uuid import uuid4

from logos_adk.checkpoint_json import JSONObject, ensure_json_object
from logos_adk.persistence import (
    get_postgresql_schema_version,
    set_postgresql_schema_version,
)
from logos_adk.postgres_utils import (
    import_psycopg,
    pooled_connect,
    resolve_database_url,
    safe_identifier,
)

from .base import (
    EventBus,
    _clone_record,
    _iso_after,
    _utcnow,
    _utcnow_iso,
    _validate_topic,
    _validate_lease_owner,
)
from .models import EVENT_BUS_COLUMNS, EVENT_BUS_SCHEMA_KIND, EVENT_BUS_SCHEMA_VERSION, EventRecord

if TYPE_CHECKING:
    pass

_logger = logging.getLogger(__name__)


def _row_to_record(row: Any) -> EventRecord:
    """Convert database row to EventRecord."""

    # Handle both dict-like and tuple-like rows
    def get_value(name: str, index: int) -> Any:
        try:
            return row[name]
        except (KeyError, TypeError):
            return row[index] if isinstance(row, (list, tuple)) else None

    return EventRecord(
        id=str(get_value("id", 0)),
        topic=str(get_value("topic", 1)),
        payload=json.loads(get_value("payload_json", 2))
        if isinstance(get_value("payload_json", 2), str)
        else get_value("payload_json", 2),
        status=str(get_value("status", 3)),
        created_at=str(get_value("created_at", 4)),
        updated_at=str(get_value("updated_at", 5)),
        available_at=str(get_value("available_at", 6)),
        campaign_id=get_value("campaign_id", 7),
        correlation_id=get_value("correlation_id", 8),
        run_id=get_value("run_id", 9),
        publisher=get_value("publisher", 10),
        lease_owner=get_value("lease_owner", 11),
        lease_expires_at=get_value("lease_expires_at", 12),
        attempt=int(get_value("attempt", 13)),
        dedupe_key=get_value("dedupe_key", 14),
        last_error=get_value("last_error", 15),
    )


class PostgreSQLEventBus(EventBus):
    """PostgreSQL-backed event bus."""

    def __init__(
        self,
        dsn: str | None = None,
        table_name: str = "events",
        *,
        max_attempts: int | None = None,
        max_queue_depth: int = 0,
    ) -> None:
        """Initialize PostgreSQL event bus.

        Args:
            dsn: Database connection string
            table_name: Table name for events
            max_attempts: Maximum retry attempts before dead-lettering
            max_queue_depth: Maximum queue depth (0 = unlimited)
        """
        super().__init__(max_attempts=max_attempts, max_queue_depth=max_queue_depth)
        self.dsn = resolve_database_url(dsn)
        if not self.dsn:
            raise ValueError(
                "PostgreSQL event bus requires a DSN or DATABASE_URL/LOGOS_ADK_DATABASE_URL"
            )
        self.table_name = safe_identifier(table_name)
        self._notify_channel = f"{self.table_name}_notify"
        self._async_listen_conn: Any | None = None
        self._listen_lock: asyncio.Lock | None = None
        self._schema_ready = False

    def _connect(self) -> Any:
        """Create database connection."""
        try:
            return pooled_connect(self.dsn)
        except Exception as exc:
            from logos_adk.errors import StateBackendError

            raise StateBackendError(f"Failed to open PostgreSQL event bus: {exc}") from exc

    def _ensure_schema(self, connection: Any) -> None:
        """Ensure database schema is up to date."""
        if self._schema_ready:
            return
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {self.table_name} (
                    id TEXT PRIMARY KEY,
                    topic TEXT NOT NULL,
                    payload_json JSONB NOT NULL,
                    status TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL,
                    available_at TEXT NOT NULL,
                    campaign_id TEXT,
                    correlation_id TEXT,
                    run_id TEXT,
                    publisher TEXT,
                    lease_owner TEXT,
                    lease_expires_at TEXT,
                    attempt INTEGER NOT NULL DEFAULT 0,
                    dedupe_key TEXT,
                    last_error TEXT
                )
                """
            )
            cursor.execute(
                f"""
                CREATE INDEX IF NOT EXISTS idx_{self.table_name}_claimable
                ON {self.table_name} (topic, status, available_at, lease_expires_at, created_at)
                """
            )
            cursor.execute(
                f"""
                CREATE INDEX IF NOT EXISTS idx_{self.table_name}_correlation
                ON {self.table_name} (correlation_id, topic, created_at)
                """
            )
            cursor.execute(
                f"""
                CREATE INDEX IF NOT EXISTS idx_{self.table_name}_campaign
                ON {self.table_name} (campaign_id, topic, created_at)
                """
            )
        recorded_version = get_postgresql_schema_version(connection, EVENT_BUS_SCHEMA_KIND)
        if recorded_version != EVENT_BUS_SCHEMA_VERSION:
            set_postgresql_schema_version(
                connection, EVENT_BUS_SCHEMA_KIND, EVENT_BUS_SCHEMA_VERSION
            )
        connection.commit()
        self._schema_ready = True

    def publish(
        self,
        topic: str,
        payload: JSONObject,
        *,
        campaign_id: str | None = None,
        correlation_id: str | None = None,
        run_id: str | None = None,
        publisher: str | None = None,
        dedupe_key: str | None = None,
        available_in_seconds: float = 0.0,
    ) -> EventRecord:
        """Publish an event to the bus."""
        self._check_backpressure()
        topic_name = _validate_topic(topic)
        created_at = _utcnow_iso()
        record = EventRecord(
            id=uuid4().hex,
            topic=topic_name,
            payload=ensure_json_object(payload, path=f"event[{topic_name}].payload"),
            status="queued",
            created_at=created_at,
            updated_at=created_at,
            available_at=_iso_after(available_in_seconds),
            campaign_id=campaign_id,
            correlation_id=correlation_id,
            run_id=run_id,
            publisher=publisher,
            dedupe_key=dedupe_key,
        )
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"""
                    INSERT INTO {self.table_name} (
                        id, topic, payload_json, status, created_at, updated_at, available_at,
                        campaign_id, correlation_id, run_id, publisher, lease_owner, lease_expires_at,
                        attempt, dedupe_key, last_error
                    ) VALUES (%s, %s, %s::jsonb, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                    """,
                    (
                        record.id,
                        record.topic,
                        json.dumps(record.payload, ensure_ascii=False, allow_nan=False),
                        record.status,
                        record.created_at,
                        record.updated_at,
                        record.available_at,
                        record.campaign_id,
                        record.correlation_id,
                        record.run_id,
                        record.publisher,
                        record.lease_owner,
                        record.lease_expires_at,
                        record.attempt,
                        record.dedupe_key,
                        record.last_error,
                    ),
                )
                # LISTEN/NOTIFY: wake up any waiting consumers immediately
                cursor.execute(
                    f"NOTIFY {self._notify_channel}, '{record.topic}'",
                )
            connection.commit()
        self._notify_observers(record)
        return _clone_record(record)

    @property
    def notify_channel(self) -> str:
        """Return the PostgreSQL NOTIFY channel name used by this bus."""
        return self._notify_channel

    async def _get_listen_conn(self) -> Any:
        """Return a persistent async psycopg connection for LISTEN."""
        if self._listen_lock is None:
            self._listen_lock = asyncio.Lock()

        async with self._listen_lock:
            if self._async_listen_conn is not None:
                if self._async_listen_conn.closed:
                    self._async_listen_conn = None
                else:
                    return self._async_listen_conn

            psycopg_mod = import_psycopg()
            conn = await psycopg_mod.AsyncConnection.connect(
                self.dsn,
                autocommit=True,
            )
            await conn.execute(f"LISTEN {self._notify_channel}")
            self._async_listen_conn = conn
            return conn

    async def wait_for_notify(self, timeout: float = 5.0) -> str | None:
        """Block until a NOTIFY arrives or timeout seconds elapse."""
        try:
            conn = await self._get_listen_conn()
        except Exception:
            return None

        try:
            async for notify in conn.notifies(timeout=timeout):
                return notify.payload
        except asyncio.TimeoutError:
            return None
        except Exception:
            try:
                await conn.close()
            except Exception:
                pass
            self._async_listen_conn = None
            return None
        return None

    async def close_listen(self) -> None:
        """Close the persistent LISTEN connection if open."""
        conn = self._async_listen_conn
        self._async_listen_conn = None
        if conn is not None:
            try:
                await conn.close()
            except Exception:
                pass

    def claim(
        self,
        consumer: str,
        topics: Iterable[str],
        *,
        lease_seconds: float = 60.0,
        correlation_id: str | None = None,
        run_id: str | None = None,
    ) -> EventRecord | None:
        """Claim an event for processing."""
        topic_names = [_validate_topic(topic) for topic in topics]
        if not topic_names:
            return None
        now = _utcnow_iso()
        lease_expires_at = _iso_after(lease_seconds)
        where_parts = [
            "topic = ANY(%s)",
            "((status = 'queued' AND available_at <= %s) OR (status = 'claimed' AND lease_expires_at IS NOT NULL AND lease_expires_at <= %s))",
        ]
        params: list[Any] = [topic_names, now, now]
        if correlation_id is not None:
            where_parts.append("correlation_id = %s")
            params.append(correlation_id)
        if run_id is not None:
            where_parts.append("run_id = %s")
            params.append(run_id)
        where_sql = " AND ".join(where_parts)

        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                # Atomic claim: UPDATE ... FROM (SELECT ... FOR UPDATE SKIP LOCKED)
                claim_params = [now, consumer, lease_expires_at, *params]
                cursor.execute(
                    f"""
                    UPDATE {self.table_name} AS t
                    SET status = 'claimed',
                        updated_at = %s,
                        lease_owner = %s,
                        lease_expires_at = %s,
                        attempt = t.attempt + 1,
                        last_error = NULL
                    FROM (
                        SELECT id FROM {self.table_name}
                        WHERE {where_sql}
                        ORDER BY available_at ASC, created_at ASC, id ASC
                        LIMIT 1
                        FOR UPDATE SKIP LOCKED
                    ) AS candidate
                    WHERE t.id = candidate.id
                    RETURNING {", ".join(f"t.{c}" for c in EVENT_BUS_COLUMNS)}
                    """,
                    tuple(claim_params),
                )
                claimed_row = cursor.fetchone()
            connection.commit()
        return None if claimed_row is None else _row_to_record(claimed_row)

    def ack(self, event_id: str, consumer: str) -> None:
        """Acknowledge event completion."""
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"SELECT {', '.join(EVENT_BUS_COLUMNS)} FROM {self.table_name} WHERE id = %s",
                    (event_id,),
                )
                row = cursor.fetchone()
                if row is None:
                    return
                _validate_lease_owner(_row_to_record(row), consumer)
                cursor.execute(
                    f"""
                    UPDATE {self.table_name}
                    SET status = 'completed',
                        updated_at = %s,
                        lease_owner = NULL,
                        lease_expires_at = NULL
                    WHERE id = %s
                    """,
                    (_utcnow_iso(), event_id),
                )
            connection.commit()

    def requeue(
        self,
        event_id: str,
        consumer: str,
        *,
        delay_seconds: float = 0.0,
        error: str | None = None,
    ) -> None:
        """Requeue event for retry."""
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"SELECT {', '.join(EVENT_BUS_COLUMNS)} FROM {self.table_name} WHERE id = %s",
                    (event_id,),
                )
                row = cursor.fetchone()
                if row is None:
                    return
                record = _row_to_record(row)
                _validate_lease_owner(record, consumer)
                # Dead Letter Queue: if max_attempts exceeded, move to dead_letter
                if self.max_attempts > 0 and record.attempt >= self.max_attempts:
                    cursor.execute(
                        f"""
                        UPDATE {self.table_name}
                        SET status = 'dead_letter',
                            updated_at = %s,
                            lease_owner = NULL,
                            lease_expires_at = NULL,
                            last_error = %s
                        WHERE id = %s
                        """,
                        (_utcnow_iso(), error, event_id),
                    )
                    connection.commit()
                    _logger.warning(
                        "Event '%s' (topic=%s) moved to dead letter queue after %d attempts",
                        event_id,
                        record.topic,
                        record.attempt,
                    )
                    return
                cursor.execute(
                    f"""
                    UPDATE {self.table_name}
                    SET status = 'queued',
                        updated_at = %s,
                        available_at = %s,
                        lease_owner = NULL,
                        lease_expires_at = NULL,
                        last_error = %s
                    WHERE id = %s
                    """,
                    (_utcnow_iso(), _iso_after(delay_seconds), error, event_id),
                )
            connection.commit()

    def requeue_dead_letter(self, event_id: str, *, delay_seconds: float = 0.0) -> bool:
        """Requeue a dead-lettered event."""
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"SELECT {', '.join(EVENT_BUS_COLUMNS)} FROM {self.table_name} WHERE id = %s",
                    (event_id,),
                )
                row = cursor.fetchone()
                if row is None:
                    return False
                record = _row_to_record(row)
                if record.status != "dead_letter":
                    return False
                cursor.execute(
                    f"""
                    UPDATE {self.table_name}
                    SET status = 'queued',
                        attempt = 0,
                        updated_at = %s,
                        available_at = %s,
                        lease_owner = NULL,
                        lease_expires_at = NULL,
                        last_error = NULL
                    WHERE id = %s
                    """,
                    (_utcnow_iso(), _iso_after(delay_seconds), event_id),
                )
            connection.commit()
        return True

    def purge_completed(self, *, max_age_seconds: float = 86400.0) -> int:
        """Purge old completed events."""
        cutoff = (_utcnow() - timedelta(seconds=max_age_seconds)).isoformat()
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"DELETE FROM {self.table_name} WHERE status IN ('completed', 'dead_letter') AND updated_at < %s",
                    (cutoff,),
                )
                count = cursor.rowcount
            connection.commit()
        return count

    def get_event(self, event_id: str) -> EventRecord | None:
        """Get event by ID."""
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"SELECT {', '.join(EVENT_BUS_COLUMNS)} FROM {self.table_name} WHERE id = %s",
                    (event_id,),
                )
                row = cursor.fetchone()
        return None if row is None else _row_to_record(row)

    def list_events(
        self,
        *,
        topic: str | None = None,
        status: str | None = None,
        campaign_id: str | None = None,
        correlation_id: str | None = None,
        run_id: str | None = None,
    ) -> list[EventRecord]:
        """List events with filters."""
        where_parts: list[str] = []
        params: list[Any] = []
        if topic is not None:
            where_parts.append("topic = %s")
            params.append(_validate_topic(topic))
        if status is not None:
            where_parts.append("status = %s")
            params.append(status)
        if campaign_id is not None:
            where_parts.append("campaign_id = %s")
            params.append(campaign_id)
        if correlation_id is not None:
            where_parts.append("correlation_id = %s")
            params.append(correlation_id)
        if run_id is not None:
            where_parts.append("run_id = %s")
            params.append(run_id)
        where_sql = "" if not where_parts else f"WHERE {' AND '.join(where_parts)}"
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"""
                    SELECT {", ".join(EVENT_BUS_COLUMNS)}
                    FROM {self.table_name}
                    {where_sql}
                    ORDER BY created_at ASC, id ASC
                    """,
                    tuple(params),
                )
                rows = cursor.fetchall()
        return [_row_to_record(row) for row in rows]


__all__ = ["PostgreSQLEventBus", "_row_to_record"]
