from __future__ import annotations

import base64
import json
from collections.abc import Sequence
from datetime import datetime, timedelta
from typing import Any, Self, TypedDict

import sqlalchemy as sa
from fansifter_common.adapters.db.models import Query
from fansifter_common.utils import timezone
from pydantic import BaseModel, ConfigDict
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.orm import (
    InstrumentedAttribute,
    Mapped,
    Session,
    aliased,
    mapped_column,
)

from resonance_engine.adapters.db import Model
from resonance_engine.dsp.enums import DSPClientName, DSPId, DSPResource
from resonance_engine.dsp.models import DSPClient
from resonance_engine.fandata.enums import (
    FanCollectionError,
    FanConnectionStatus,
)
from resonance_engine.fandata.types import (
    ActivityBucket,
    CollectionGranularity,
    CollectionStats,
    ConnectionStats,
    FanView,
    FanViewPaginated,
)
from resonance_engine.fandata.utils import make_fan_id

# ─── Fan connection ──────────────────────────────────────────────────────────


def _explain_rows(session: Session, stmt: sa.Select[Any]) -> int:
    """Approximate row count for `stmt` via EXPLAIN (planner estimate, no scan;
    DSQL keeps stats fresh with background ANALYZE). COUNT(*) would scan."""
    sql = stmt.compile(
        dialect=session.get_bind().dialect,
        compile_kwargs={"literal_binds": True},
    )
    plan = session.execute(sa.text(f"EXPLAIN (FORMAT JSON) {sql}")).scalar()
    if isinstance(plan, str):
        plan = json.loads(plan)
    if not plan:
        return 0
    return max(0, int(plan[0]["Plan"]["Plan Rows"]))


class FanConnectionFilter(BaseModel):
    model_config = ConfigDict(frozen=True)

    dsp_client_id: int | None = None
    status: FanConnectionStatus | None = FanConnectionStatus.active
    not_collected_since: datetime | None = None
    not_last_dispatched_since: datetime | None = None
    max_consecutive_failures: int | None = None
    limit: int | None = None

    def apply[S: tuple[Any, ...]](self, stmt: sa.Select[S]) -> sa.Select[S]:
        if self.dsp_client_id is not None:
            stmt = stmt.where(FanConnection.dsp_client_id == self.dsp_client_id)

        if self.status is not None:
            stmt = stmt.where(FanConnection.status == self.status)

        if self.not_collected_since is not None:
            stmt = stmt.where(
                sa.or_(
                    # NULL = collect ASAP (new connection or re-connected).
                    FanConnection.last_collected_at.is_(None),
                    FanConnection.last_collected_at < self.not_collected_since,
                )
            )

        if self.not_last_dispatched_since is not None:
            stmt = stmt.where(
                sa.or_(
                    # Never dispatched — always eligible.
                    FanConnection.last_dispatched_at.is_(None),
                    # Collected after last dispatch → Lambda completed successfully.
                    FanConnection.last_collected_at > FanConnection.last_dispatched_at,
                    # Safety net: dispatch lock expired (covers DLQ / Lambda crash).
                    FanConnection.last_dispatched_at < self.not_last_dispatched_since,
                )
            )

        if self.max_consecutive_failures is not None:
            stmt = stmt.where(
                sa.or_(
                    FanConnection.consecutive_failures.is_(None),
                    FanConnection.consecutive_failures < self.max_consecutive_failures,
                )
            )

        if self.limit is not None:
            stmt = stmt.limit(self.limit)

        return stmt


class FanConnectionRow(TypedDict):
    fan_id: str
    dsp_id: DSPId
    dsp_client_id: int
    token_encrypted: str


class FanConnectionViewFilter(BaseModel):
    model_config = ConfigDict(frozen=True)

    dsp_id: DSPId | None = None
    dsp_client_id: int | None = None
    dsp_client_names: tuple[DSPClientName, ...] | None = None
    status: FanConnectionStatus | None = None
    search: str | None = None
    is_collected: bool | None = None


def _encode_fan_view_cursor(fan_id: str, dsp_id: str) -> str:
    payload = json.dumps([fan_id, dsp_id]).encode()
    return base64.urlsafe_b64encode(payload).decode().rstrip("=")


def _decode_fan_view_cursor(cursor: str) -> tuple[str, str]:
    padded = cursor + "=" * (-len(cursor) % 4)
    fan_id, dsp_id = json.loads(base64.urlsafe_b64decode(padded))
    return str(fan_id), str(dsp_id)


class FanConnectionQuery(Query["FanConnection"]):
    def filter(self, f: FanConnectionFilter) -> Self:
        self._query = f.apply(sa.select(FanConnection))
        return self

    def approx_stats(self) -> ConnectionStats:
        """Approximate total/active/revoked over the chained query via planner
        estimates (no table scan)."""
        total = _explain_rows(self.session, self._query)
        active = _explain_rows(
            self.session,
            self._query.where(FanConnection.status == FanConnectionStatus.active),
        )
        return ConnectionStats(
            total=total,
            active=min(active, total),
            revoked=max(0, total - active),
        )

    def connect(self, rows: list[FanConnectionRow]) -> None:
        """Insert or refresh fan connections, queuing each for the next collect."""
        now = timezone.now()
        # connect_rank = -epoch so a plain ASC index scan returns newest-connect
        # first (DSQL has no DESC index keys nor backward scans).
        rank = -int(now.timestamp())
        stmt = insert(FanConnection)
        stmt = stmt.on_conflict_do_update(
            index_elements=["fan_id", "dsp_id", "dsp_client_id"],
            set_={
                "token_encrypted": stmt.excluded.token_encrypted,
                "status": FanConnectionStatus.active,
                "token_refreshed_at": now,
                "last_collected_at": None,
                "last_dispatched_at": None,
                "consecutive_failures": None,
                "connect_rank": rank,
                "last_connected_at": now,
            },
        )
        self.session.execute(
            stmt,
            [{**row, "connect_rank": rank, "last_connected_at": now} for row in rows],
        )

    def update_token(
        self,
        *,
        fan_id: str,
        dsp_id: DSPId,
        dsp_client_id: int,
        token_encrypted: str,
    ) -> None:
        self.session.execute(
            sa.update(FanConnection)
            .where(
                FanConnection.fan_id == fan_id,
                FanConnection.dsp_id == dsp_id,
                FanConnection.dsp_client_id == dsp_client_id,
            )
            .values(
                token_encrypted=token_encrypted,
                token_refreshed_at=sa.func.now(),
            )
        )

    def order_by_collection_priority(self) -> Self:
        # Newest connect first. connect_rank = -epoch(last connect), so a plain ASC
        # index scan (ix_fan_connection_planner) returns freshest-connected first and
        # early-terminates at LIMIT — no sort, no backward scan (DSQL has neither).
        # NULL connect_rank (never connected) sorts last (ASC NULLS LAST default).
        self._query = self._query.order_by(
            FanConnection.connect_rank.asc(),
        )
        return self

    def mark_revoked(self, *, fan_id: str, dsp_id: DSPId, dsp_client_id: int) -> None:
        self.session.execute(
            sa.update(FanConnection)
            .where(
                FanConnection.fan_id == fan_id,
                FanConnection.dsp_id == dsp_id,
                FanConnection.dsp_client_id == dsp_client_id,
            )
            .values(status=FanConnectionStatus.revoked)
        )

    def view_paginate(
        self,
        f: FanConnectionViewFilter,
        *,
        cursor: str | None = None,
        limit: int = 20,
    ) -> FanViewPaginated:
        filter_conditions: list[Any] = []
        if f.dsp_id is not None:
            filter_conditions.append(FanConnection.dsp_id == f.dsp_id)
        if f.dsp_client_id is not None:
            filter_conditions.append(FanConnection.dsp_client_id == f.dsp_client_id)
        if f.dsp_client_names:
            names = set(f.dsp_client_names)
            fc, dc = aliased(FanConnection), aliased(DSPClient)
            has_all_clients = (
                sa.select(fc.fan_id)
                .join(dc, fc.dsp_client_id == dc.id)
                .where(dc.name.in_(names))
                .group_by(fc.fan_id)
                .having(sa.func.count(sa.distinct(dc.name)) == len(names))
            )
            filter_conditions.append(FanConnection.fan_id.in_(has_all_clients))
        if f.status is not None:
            filter_conditions.append(FanConnection.status == f.status)
        if f.search:
            # search matches a raw fan_id or an email (hashed to its fan_id).
            filter_conditions.append(
                sa.or_(
                    FanConnection.fan_id == f.search,
                    FanConnection.fan_id == make_fan_id(f.search),
                )
            )
        if f.is_collected is not None:
            collected = (
                sa.exists()
                .where(
                    FanCollectionState.fan_id == FanConnection.fan_id,
                    FanCollectionState.dsp_id == FanConnection.dsp_id,
                )
                .correlate(FanConnection)
            )
            filter_conditions.append(collected if f.is_collected else ~collected)

        fan_keys_stmt = (
            sa.select(FanConnection.fan_id, FanConnection.dsp_id)
            .where(*filter_conditions)
            .distinct()
        )

        if cursor:
            cur_user_id, cur_dsp_id = _decode_fan_view_cursor(cursor)
            fan_keys_stmt = fan_keys_stmt.where(
                sa.or_(
                    FanConnection.fan_id > cur_user_id,
                    sa.and_(
                        FanConnection.fan_id == cur_user_id,
                        FanConnection.dsp_id > cur_dsp_id,
                    ),
                )
            )
        fan_keys_stmt = fan_keys_stmt.order_by(
            FanConnection.fan_id, FanConnection.dsp_id
        ).limit(limit + 1)

        fan_keys = self.session.execute(fan_keys_stmt).all()
        has_more = len(fan_keys) > limit
        page_keys = fan_keys[:limit]

        if not page_keys:
            return FanViewPaginated(items=[], next_cursor=None)

        page_filter = sa.or_(
            *[
                sa.and_(
                    FanConnection.fan_id == k.fan_id,
                    FanConnection.dsp_id == k.dsp_id,
                )
                for k in page_keys
            ]
        )

        expand_stmt = (
            sa.select(
                FanConnection.fan_id,
                FanConnection.dsp_id,
                FanConnection.status,
                FanConnection.created_at.label("first_seen_at"),
                FanCollectionState.last_collected_at,
                FanCollectionState.last_collection_error,
                sa.func.coalesce(FanCollectionState.consecutive_failures, 0).label(
                    "consecutive_failures"
                ),
                DSPClient.id.label("dsp_client_id"),
                DSPClient.name.label("dsp_client_name"),
                DSPClient.display_name.label("dsp_client_display_name"),
                FanCollectionState.profile_collected_at,
                FanCollectionState.top_artists_collected_at,
                FanCollectionState.top_tracks_collected_at,
                FanCollectionState.recently_played_collected_at,
                FanCollectionState.playlists_collected_at,
                FanCollectionState.saved_albums_collected_at,
                FanCollectionState.saved_tracks_collected_at,
                FanCollectionState.followed_artists_collected_at,
            )
            .outerjoin(
                FanCollectionState,
                sa.and_(
                    FanConnection.fan_id == FanCollectionState.fan_id,
                    FanConnection.dsp_id == FanCollectionState.dsp_id,
                ),
            )
            .join(DSPClient, FanConnection.dsp_client_id == DSPClient.id)
            .where(*filter_conditions, page_filter)
            .order_by(
                FanConnection.fan_id,
                FanConnection.dsp_id,
                FanConnection.dsp_client_id,
            )
        )
        rows = self.session.execute(expand_stmt).mappings().all()

        next_cursor = (
            _encode_fan_view_cursor(page_keys[-1].fan_id, page_keys[-1].dsp_id)
            if has_more
            else None
        )
        return FanViewPaginated(
            items=[FanView.model_validate(dict(r)) for r in rows],
            next_cursor=next_cursor,
        )


class FanConnection(Model, kw_only=True):
    __tablename__ = "fan_connection"
    __table_args__ = (
        # Covers the multi-client AND filter (view_paginate): filter by
        # dsp_client_id, group by fan_id — index-only, no heap scan.
        sa.Index(
            "ix_fan_connection_dsp_client_id_fan_id",
            "dsp_client_id",
            "fan_id",
        ),
        # created_at range scans (no live query as of now — kept for reporting).
        sa.Index("ix_fan_connection_created_at", "created_at"),
        # Planner fan-selection: forward index scan + LIMIT, no sort. connect_rank
        # (= -epoch of last connect) leads, so the scan returns newest-connect-first
        # and early-terminates. INCLUDE carries the eligibility-filter columns so the
        # not_collected_since / dispatch-lock / failure checks stay index-only.
        sa.Index(
            "ix_fan_connection_planner",
            "dsp_client_id",
            "status",
            "connect_rank",
            postgresql_include=[
                "last_collected_at",
                "last_dispatched_at",
                "consecutive_failures",
            ],
        ),
    )

    fan_id: Mapped[str] = mapped_column(primary_key=True)
    dsp_id: Mapped[DSPId] = mapped_column(primary_key=True)
    dsp_client_id: Mapped[int] = mapped_column(primary_key=True)
    token_encrypted: Mapped[str]
    token_refreshed_at: Mapped[datetime | None] = mapped_column(default=None)
    status: Mapped[FanConnectionStatus] = mapped_column(
        default=FanConnectionStatus.active
    )
    created_at: Mapped[datetime] = mapped_column(
        default_factory=timezone.now, server_default=sa.func.now()
    )
    # -epoch of last (re)connect — planner sort key (ASC = newest connect first).
    connect_rank: Mapped[int | None] = mapped_column(default=None)
    # Readable timestamp of the last (re)connect; connect_rank is its sort key.
    last_connected_at: Mapped[datetime | None] = mapped_column(default=None)
    # Denormalized from fan_collection_state for index-only planner selection.
    last_collected_at: Mapped[datetime | None] = mapped_column(default=None)
    last_dispatched_at: Mapped[datetime | None] = mapped_column(default=None)
    consecutive_failures: Mapped[int | None] = mapped_column(default=None)

    query = FanConnectionQuery.as_descriptor()


# ─── Fan collection state ─────────────────────────────────────────────────────


class FanCollectionStateQuery(Query["FanCollectionState"]):
    def approx_stats(self) -> CollectionStats:
        """Approximate total/healthy/with_errors over the chained query via planner
        estimates (no scan)."""
        total = _explain_rows(self.session, self._query)
        healthy = _explain_rows(
            self.session,
            self._query.where(FanCollectionState.last_collection_error.is_(None)),
        )
        return CollectionStats(
            total=total,
            healthy=min(healthy, total),
            with_errors=max(0, total - healthy),
        )

    def bulk_get(
        self,
        *,
        fan_ids: list[str],
        dsp_id: DSPId,
    ) -> dict[str, FanCollectionState]:
        """Fetch FanCollectionState for multiple fans in one query.

        Returns a mapping of fan_id → FanCollectionState for rows that
        exist; fans with no state record are absent from the result.
        """
        if not fan_ids:
            return {}
        rows = self.session.execute(
            sa.select(FanCollectionState).where(
                FanCollectionState.fan_id.in_(fan_ids),
                FanCollectionState.dsp_id == dsp_id,
            )
        ).scalars()
        return {row.fan_id: row for row in rows}

    def activity(
        self,
        *,
        days: int = 7,
        granularity: CollectionGranularity = "daily",
    ) -> list[ActivityBucket]:
        cutoff = timezone.now() - timedelta(days=days)
        bucket_expr = (
            sa.func.date_trunc("hour", FanCollectionState.last_collected_at)
            if granularity == "hourly"
            else sa.func.date(FanCollectionState.last_collected_at)
        )

        stmt = (
            sa.select(
                bucket_expr.label("bucket"),
                FanCollectionState.dsp_id,
                sa.func.count()
                .filter(FanCollectionState.last_collection_error.is_(None))
                .label("success"),
                sa.func.count()
                .filter(
                    FanCollectionState.last_collection_error.in_(
                        (
                            FanCollectionError.token_error,
                            FanCollectionError.api_error,
                        )
                    )
                )
                .label("errors"),
                sa.func.count()
                .filter(
                    FanCollectionState.last_collection_error
                    == FanCollectionError.rate_limited
                )
                .label("throttled"),
            )
            .where(FanCollectionState.last_collected_at >= cutoff)
            .group_by(bucket_expr, FanCollectionState.dsp_id)
            .order_by(bucket_expr, FanCollectionState.dsp_id)
        )

        rows = self.session.execute(stmt).all()
        return [
            ActivityBucket(
                bucket=str(row.bucket),
                dsp_id=row.dsp_id,
                success=row.success,
                errors=row.errors,
                throttled=row.throttled,
            )
            for row in rows
        ]

    def stamp_dispatched(
        self, pairs: Sequence[tuple[str, DSPId]], *, now: datetime
    ) -> None:
        """Stamp last_dispatched_at = now on the fan_connection rows the planner
        reads (all client-rows of each given (fan_id, dsp_id) pair). Plain tuples
        (not ORM objects) so the planner can run in its own transaction."""
        if not pairs:
            return
        self.session.execute(
            sa.update(FanConnection)
            .where(sa.tuple_(FanConnection.fan_id, FanConnection.dsp_id).in_(pairs))
            .values(last_dispatched_at=now)
        )

    def record(
        self,
        *,
        fan_id: str,
        dsp_id: DSPId,
        dsp_client_id: int,
        collected_at: datetime,
        error: FanCollectionError | None,
        resources: dict[DSPResource, datetime] | None = None,
    ) -> None:
        # rate_limited is DSP-level (not fan-specific); any other error counts as fan failure.
        is_failure = error is not None and error != FanCollectionError.rate_limited
        resource_values = {
            _RESOURCE_COLUMN[s].key: ts for s, ts in (resources or {}).items()
        }
        stmt = insert(FanCollectionState).values(
            fan_id=fan_id,
            dsp_id=dsp_id,
            last_dsp_client_id=dsp_client_id,
            last_collected_at=collected_at,
            last_collection_error=error,
            consecutive_failures=1 if is_failure else 0,
            **resource_values,
        )
        if is_failure:
            new_failures = FanCollectionState.consecutive_failures + 1
        elif error is None:
            new_failures = 0
        else:
            # rate_limited — DSP-level signal, not a fan-specific failure. Leave counter alone.
            new_failures = FanCollectionState.consecutive_failures
        stmt = stmt.on_conflict_do_update(
            index_elements=["fan_id", "dsp_id"],
            set_={
                "last_dsp_client_id": stmt.excluded.last_dsp_client_id,
                "last_collected_at": stmt.excluded.last_collected_at,
                "last_collection_error": stmt.excluded.last_collection_error,
                "consecutive_failures": new_failures,
                **resource_values,
            },
        )
        self.session.execute(stmt)
        # Mirror onto the fan's connection rows (denormalized planner columns);
        # correlated from fcs so server-computed consecutive_failures matches.
        # last_collected_at moving forward pushes the fan back by staleness.
        # Skip if last_dispatched_at IS NULL: connect() clears it on re-connect, so
        # a collect finishing after a re-connect must not overwrite the fresh re-queue.
        self.session.execute(
            sa.update(FanConnection)
            .values(
                last_collected_at=FanCollectionState.last_collected_at,
                consecutive_failures=FanCollectionState.consecutive_failures,
            )
            .where(
                FanConnection.fan_id == FanCollectionState.fan_id,
                FanConnection.dsp_id == FanCollectionState.dsp_id,
                FanConnection.fan_id == fan_id,
                FanConnection.dsp_id == dsp_id,
                FanConnection.last_dispatched_at.is_not(None),
            )
        )


class FanCollectionState(Model, kw_only=True):
    __tablename__ = "fan_collection_state"

    fan_id: Mapped[str] = mapped_column(primary_key=True)
    dsp_id: Mapped[DSPId] = mapped_column(primary_key=True)
    last_dsp_client_id: Mapped[int]
    last_collected_at: Mapped[datetime] = mapped_column(index=True)
    last_collection_error: Mapped[FanCollectionError | None] = mapped_column(
        default=None
    )
    consecutive_failures: Mapped[int] = mapped_column(default=0)
    # TODO: remove it later
    # last_dispatched_at: Mapped[datetime | None] = mapped_column(default=None)
    profile_collected_at: Mapped[datetime | None] = mapped_column(default=None)
    top_artists_collected_at: Mapped[datetime | None] = mapped_column(default=None)
    top_tracks_collected_at: Mapped[datetime | None] = mapped_column(default=None)
    recently_played_collected_at: Mapped[datetime | None] = mapped_column(default=None)
    playlists_collected_at: Mapped[datetime | None] = mapped_column(default=None)
    saved_albums_collected_at: Mapped[datetime | None] = mapped_column(default=None)
    saved_tracks_collected_at: Mapped[datetime | None] = mapped_column(default=None)
    followed_artists_collected_at: Mapped[datetime | None] = mapped_column(default=None)

    query = FanCollectionStateQuery.as_descriptor()


_RESOURCE_COLUMN: dict[DSPResource, InstrumentedAttribute[datetime | None]] = {
    DSPResource.profile: FanCollectionState.profile_collected_at,
    DSPResource.top_artists: FanCollectionState.top_artists_collected_at,
    DSPResource.top_tracks: FanCollectionState.top_tracks_collected_at,
    DSPResource.recently_played: FanCollectionState.recently_played_collected_at,
    DSPResource.playlists: FanCollectionState.playlists_collected_at,
    DSPResource.saved_albums: FanCollectionState.saved_albums_collected_at,
    DSPResource.saved_tracks: FanCollectionState.saved_tracks_collected_at,
    DSPResource.followed_artists: FanCollectionState.followed_artists_collected_at,
}

if set(_RESOURCE_COLUMN) != set(DSPResource):
    raise AssertionError("_RESOURCE_COLUMN must cover every DSPResource value")