from __future__ import annotations

import base64
import json
import uuid
from collections import defaultdict
from datetime import datetime, timedelta
from statistics import median
from typing import cast

from fansifter_common.adapters.db.models import Query
from fansifter_common.adapters.db.types import PydanticType
from fansifter_common.utils import timezone
from sqlalchemy import Index, and_, case, exists, func, or_, select, update
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.engine import CursorResult
from sqlalchemy.orm import Mapped, mapped_column

from resonance_engine.adapters.db import Model, db
from resonance_engine.dsp.enums import DSPClientName
from resonance_engine.dsp.models import DSPClient
from resonance_engine.fandata.models import FanConnectionFilter
from resonance_engine.fandata.types import CollectionGranularity
from resonance_engine.tasks.enums import FanoutSource, TaskFinishedReason, TaskStatus
from resonance_engine.tasks.types import (
    CollectActivityBucket,
    CollectAggStats,
    FanoutAggStats,
    FanoutSignals,
    FanoutTaskView,
    FanoutTaskViewPage,
)


def _encode_fanout_cursor(started_at: datetime, task_id: uuid.UUID) -> str:
    payload = json.dumps([started_at.isoformat(), str(task_id)]).encode()
    return base64.urlsafe_b64encode(payload).decode().rstrip("=")


def _decode_fanout_cursor(cursor: str) -> tuple[datetime, uuid.UUID]:
    padded = cursor + "=" * (-len(cursor) % 4)
    started_at_iso, task_id_str = json.loads(base64.urlsafe_b64decode(padded))
    return datetime.fromisoformat(started_at_iso), uuid.UUID(task_id_str)


# ─── Fanout task ─────────────────────────────────────────────────────────────


class FanoutTaskQuery(Query["FanoutTask"]):
    def last_scheduled_started_at(self) -> datetime | None:
        """started_at of the most recent scheduled run — anchor for estimating
        the next cron tick. None when no scheduled run has happened yet."""
        return db.session.scalar(
            select(FanoutTask.started_at)
            .where(FanoutTask.source == FanoutSource.scheduled)
            .order_by(FanoutTask.started_at.desc(), FanoutTask.id.desc())
            .limit(1)
        )

    def update_dispatch(
        self,
        fanout_task_id: uuid.UUID,
        *,
        fans_dispatched: int,
        messages_sent: int,
    ) -> None:
        """Set dispatch counters after async (SQS) fanout completes."""
        db.session.execute(
            update(FanoutTask)
            .where(FanoutTask.id == fanout_task_id)
            .values(
                fans_dispatched=fans_dispatched,
                messages_sent=messages_sent,
            )
        )

    def increment_dispatch(
        self,
        fanout_task_id: uuid.UUID,
        *,
        fans: int,
        messages: int,
    ) -> None:
        """Atomically increment dispatch counters after each inline (sync) collect batch."""
        db.session.execute(
            update(FanoutTask)
            .where(FanoutTask.id == fanout_task_id)
            .values(
                fans_dispatched=FanoutTask.fans_dispatched + fans,
                messages_sent=FanoutTask.messages_sent + messages,
            )
        )

    def paginate(
        self, *, cursor: str | None = None, limit: int = 20
    ) -> FanoutTaskViewPage:
        """Cursor-paginated page, newest first; each row carries collect
        aggregates (collected/errors/requests/throttled) from one query."""
        total = (
            db.session.scalar(select(func.count()).select_from(self._query.subquery()))
            or 0
        )
        collect = (
            select(
                CollectTask.fanout_task_id.label("fanout_id"),
                func.sum(CollectTask.fans_processed).label("collected"),
                func.sum(CollectTask.fans_errors).label("errors"),
                func.sum(CollectTask.fans_stale_tokens).label("stale"),
                func.sum(CollectTask.fans_skipped).label("skipped"),
                func.sum(CollectTask.requests).label("requests"),
                func.sum(CollectTask.requests_rate_limited).label("throttled"),
                func.count(CollectTask.finished_at).label("done"),
            )
            .group_by(CollectTask.fanout_task_id)
            .subquery()
        )
        q = (
            self._query.outerjoin(collect, collect.c.fanout_id == FanoutTask.id)
            .add_columns(
                func.coalesce(collect.c.collected, 0),
                func.coalesce(collect.c.errors, 0),
                func.coalesce(collect.c.stale, 0),
                func.coalesce(collect.c.skipped, 0),
                func.coalesce(collect.c.requests, 0),
                func.coalesce(collect.c.throttled, 0),
                func.coalesce(collect.c.done, 0),
            )
            .order_by(FanoutTask.started_at.desc(), FanoutTask.id.desc())
        )
        if cursor:
            cursor_dt, cursor_id = _decode_fanout_cursor(cursor)
            q = q.where(
                or_(
                    FanoutTask.started_at < cursor_dt,
                    and_(
                        FanoutTask.started_at == cursor_dt,
                        FanoutTask.id < cursor_id,
                    ),
                )
            )
        rows = db.session.execute(q.limit(limit + 1)).all()
        has_more = len(rows) > limit
        rows = rows[:limit]
        items = [
            FanoutTaskView(
                id=task.id,
                started_at=task.started_at,
                finished_at=task.finished_at,
                source=task.source,
                status=task.status,
                finished_reason=task.finished_reason,
                fans_dispatched=task.fans_dispatched,
                messages_sent=task.messages_sent,
                filters=task.filters,
                fans_collected=int(collected),
                fans_errors=int(errors),
                fans_stale_tokens=int(stale),
                fans_skipped=int(skipped),
                requests=int(requests),
                requests_rate_limited=int(throttled),
                collect_tasks_done=int(done),
            )
            for task, collected, errors, stale, skipped, requests, throttled, done in rows
        ]
        next_cursor = (
            _encode_fanout_cursor(rows[-1][0].started_at, rows[-1][0].id)
            if has_more and rows
            else None
        )
        return FanoutTaskViewPage(items=items, next_cursor=next_cursor, total=total)

    def aggregate_stats(self) -> FanoutAggStats:
        """Return aggregate stats across all fanout tasks."""
        row = db.session.execute(
            select(
                func.count(FanoutTask.id).label("total"),
                func.count(FanoutTask.id)
                .filter(FanoutTask.status == TaskStatus.running)
                .label("running"),
                func.max(FanoutTask.started_at).label("last_started_at"),
            )
        ).one()
        last_fans = db.session.scalar(
            select(FanoutTask.fans_dispatched)
            .order_by(FanoutTask.started_at.desc(), FanoutTask.id.desc())
            .limit(1)
        )
        return FanoutAggStats(
            total=row.total,
            running=row.running,
            last_started_at=row.last_started_at,
            last_fans_dispatched=last_fans,
        )

    def close_stale(self, *, dispatch_stale_s: int) -> int:
        """Close fanouts stuck before dispatch finished (messages_sent == 0) —
        safe to time out since no collect Lambdas ran. Dispatched fanouts are
        excluded (SQS retries outlast any window); those close via
        close_if_all_collected or the DLQ Lambda. Returns the count closed."""
        now = timezone.now()
        cutoff = now - timedelta(seconds=dispatch_stale_s)

        stale_ids = select(FanoutTask.id).where(
            FanoutTask.status == TaskStatus.running,
            FanoutTask.messages_sent == 0,
            FanoutTask.started_at < cutoff,
        )
        db.session.execute(
            update(CollectTask)
            .where(
                CollectTask.fanout_task_id.in_(stale_ids),
                CollectTask.status == TaskStatus.running,
            )
            .values(
                status=TaskStatus.stale,
                finished_at=now,
                finished_reason=TaskFinishedReason.timed_out,
            )
        )
        result = db.session.execute(
            update(FanoutTask)
            .where(FanoutTask.id.in_(stale_ids))
            .values(
                status=TaskStatus.done,
                finished_at=now,
                finished_reason=TaskFinishedReason.timed_out,
            )
        )
        return cast(CursorResult, result).rowcount

    def close_if_all_collected(self, fanout_task_id: uuid.UUID) -> bool:
        """Atomically close the fanout task when all dispatched batches have settled.

        Guards against a race where a collect Lambda finishes and calls this method
        while other batches are still queued in SQS but have not yet created their
        CollectTask rows (status=running). Without the count check, the ~still_running
        predicate would fire prematurely — no running rows yet ≠ all batches done.

        Closing condition:
          - messages_sent > 0 (dispatch has committed — guards against pre-commit window)
          - count(terminal CollectTasks) == messages_sent (every batch has settled)

        finished_reason is timed_out if any sibling ended as stale (DLQ path),
        done if all siblings completed successfully.
        """
        now = timezone.now()

        # Count collect tasks that have reached a terminal state (not running).
        terminal_count = (
            select(func.count())
            .where(
                CollectTask.fanout_task_id == fanout_task_id,
                CollectTask.status != TaskStatus.running,
            )
            .scalar_subquery()
        )
        any_stale = exists(
            select(CollectTask.id).where(
                CollectTask.fanout_task_id == fanout_task_id,
                CollectTask.status == TaskStatus.stale,
            )
        )
        finished_reason = case(
            (any_stale, TaskFinishedReason.timed_out),
            else_=TaskFinishedReason.done,
        )
        result = db.session.execute(
            update(FanoutTask)
            .where(
                FanoutTask.id == fanout_task_id,
                FanoutTask.status == TaskStatus.running,
                # All dispatched batches must have settled — catches the race where
                # SQS-queued batches have no CollectTask row yet.
                FanoutTask.messages_sent > 0,
                terminal_count == FanoutTask.messages_sent,
            )
            .values(
                status=TaskStatus.done,
                finished_at=now,
                finished_reason=finished_reason,
            )
        )
        return cast(CursorResult, result).rowcount == 1


class FanoutTask(Model, kw_only=True):
    """One fanout cron tick — covers all DSP clients dispatched in a single run."""

    __tablename__ = "task_fanout"
    __table_args__ = (Index("ix_task_fanout_started_at", "started_at"),)

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), primary_key=True, default_factory=uuid.uuid4, init=False
    )
    started_at: Mapped[datetime]
    finished_at: Mapped[datetime | None] = mapped_column(default=None)
    source: Mapped[FanoutSource]
    status: Mapped[TaskStatus]
    finished_reason: Mapped[TaskFinishedReason | None] = mapped_column(default=None)
    fans_dispatched: Mapped[int] = mapped_column(default=0)
    messages_sent: Mapped[int] = mapped_column(default=0)
    filters: Mapped[FanConnectionFilter] = mapped_column(
        PydanticType(FanConnectionFilter, exclude_unset=True),
        default_factory=FanConnectionFilter,
    )

    query = FanoutTaskQuery.as_descriptor()


# ─── Collect task ─────────────────────────────────────────────────────────────


_SIGNAL_FANOUTS = 3  # fanouts the throughput signals are medianed over (noise damping)


class CollectTaskQuery(Query["CollectTask"]):
    def aggregate_stats(
        self, dsp_client_name: DSPClientName | None = None
    ) -> CollectAggStats:
        """Return aggregate stats across collect tasks, optionally filtered by DSP client."""
        stmt = select(
            func.count(CollectTask.id).label("total"),
            func.count(CollectTask.id)
            .filter(CollectTask.status == TaskStatus.running)
            .label("running"),
            func.count(CollectTask.id)
            .filter(CollectTask.status == TaskStatus.stale)
            .label("stale"),
            func.max(CollectTask.started_at).label("last_started_at"),
        )
        if dsp_client_name is not None:
            stmt = stmt.where(CollectTask.dsp_client_name == dsp_client_name)
        row = db.session.execute(stmt).one()
        return CollectAggStats(
            total=row.total,
            running=row.running,
            stale=row.stale,
            last_started_at=row.last_started_at,
        )

    def pending_fans(self, dsp_client_name: DSPClientName) -> int:
        """Unfinished fans in this client's running fanouts — subtracted from the
        next dispatch as headroom. Scoped to running fanouts to keep it index-seekable
        (fanout_task_id) rather than scanning the client's whole collect history."""
        running = select(FanoutTask.id).where(FanoutTask.status == TaskStatus.running)
        return (
            db.session.scalar(
                select(func.coalesce(func.sum(CollectTask.fans_total), 0)).where(
                    CollectTask.fanout_task_id.in_(running),
                    CollectTask.dsp_client_name == dsp_client_name,
                    CollectTask.finished_at.is_(None),
                    CollectTask.status != TaskStatus.stale,
                )
            )
            or 0
        )

    def activity(
        self,
        *,
        days: int = 7,
        granularity: CollectionGranularity = "daily",
        dsp_client_name: DSPClientName | None = None,
    ) -> list[CollectActivityBucket]:
        """Aggregate processed/errors/rate_limited from completed collect tasks."""
        cutoff = timezone.now() - timedelta(days=days)
        bucket_expr = (
            func.date_trunc("hour", CollectTask.started_at)
            if granularity == "hourly"
            else func.date(CollectTask.started_at)
        )
        filters = [
            CollectTask.started_at >= cutoff,
            CollectTask.status == TaskStatus.done,
        ]
        if dsp_client_name is not None:
            filters.append(CollectTask.dsp_client_name == dsp_client_name)
        stmt = (
            select(
                bucket_expr.label("bucket"),
                CollectTask.dsp_client_name,
                DSPClient.display_name.label("dsp_client_display_name"),
                func.sum(CollectTask.fans_processed).label("fans_processed"),
                func.sum(CollectTask.fans_errors).label("fans_errors"),
                func.sum(CollectTask.fans_stale_tokens).label("fans_stale_tokens"),
                func.sum(CollectTask.requests_rate_limited).label(
                    "requests_rate_limited"
                ),
                func.sum(CollectTask.requests).label("requests"),
            )
            .join(DSPClient, DSPClient.name == CollectTask.dsp_client_name)
            .where(*filters)
            .group_by(bucket_expr, CollectTask.dsp_client_name, DSPClient.display_name)
            .order_by(bucket_expr, CollectTask.dsp_client_name)
        )
        rows = db.session.execute(stmt).all()
        return [
            CollectActivityBucket(
                bucket=str(row.bucket),
                dsp_client_name=row.dsp_client_name,
                dsp_client_display_name=row.dsp_client_display_name,
                fans_processed=row.fans_processed or 0,
                fans_errors=row.fans_errors or 0,
                fans_stale_tokens=row.fans_stale_tokens or 0,
                requests_rate_limited=row.requests_rate_limited or 0,
                requests=row.requests or 0,
            )
            for row in rows
        ]

    def fanout_signals(
        self, dsp_client_name: DSPClientName, *, fanouts: int = _SIGNAL_FANOUTS
    ) -> FanoutSignals | None:
        """Median throughput signals over the last few fully-drained fanouts, for
        sizing the next one. Spans are the FULL drain (max finish): the slow tail
        is real concurrency-queue drain that p90 dropped, which over-sized the
        controller. `None` if there are no completed fanouts to measure."""

        # 1. Recent fully-drained fanouts (no running task) — from task_fanout.
        running = (
            select(1)
            .where(
                CollectTask.fanout_task_id == FanoutTask.id,
                CollectTask.status == TaskStatus.running,
            )
            .exists()
        )
        recent_ids = (
            db.session.execute(
                select(FanoutTask.id)
                .where(~running)
                .order_by(FanoutTask.started_at.desc())
                .limit(fanouts)
            )
            .scalars()
            .all()
        )

        if not recent_ids:
            return None

        # Finished tasks only — a running task has full fans_total but partial drain.
        where_clause = (
            CollectTask.fanout_task_id.in_(recent_ids),
            CollectTask.dsp_client_name == dsp_client_name,
            CollectTask.status != TaskStatus.stale,
            CollectTask.finished_at.isnot(None),
        )

        # 2. Counts aggregated in SQL (≤N rows, one per fanout).
        count_rows = db.session.execute(
            select(
                CollectTask.fanout_task_id,
                func.sum(CollectTask.requests),
                func.sum(CollectTask.fans_total),
                func.sum(CollectTask.requests_rate_limited),
            )
            .where(*where_clause)
            .group_by(CollectTask.fanout_task_id)
        ).all()
        counts_by_fanout = {
            fanout_id: (requests, fans, rate_limited)
            for fanout_id, requests, fans, rate_limited in count_rows
        }

        # 3. Per-task timestamps for the span, grouped by fanout in Python.
        ts_rows = db.session.execute(
            select(
                CollectTask.fanout_task_id,
                CollectTask.started_at,
                CollectTask.finished_at,
            ).where(*where_clause)
        ).all()
        ts_by_fanout = defaultdict(list)
        for row in ts_rows:
            ts_by_fanout[row.fanout_task_id].append(row)

        # 4. Combine counts + timestamps per fanout.
        rps_list, rpf_list, duration_list = [], [], []
        time_per_fan_list, fans_per_s_list = [], []
        latency_list, throttle_list = [], []

        for fanout_id, group in ts_by_fanout.items():
            counts = counts_by_fanout.get(fanout_id)
            if counts is None:
                continue
            requests, fans, rate_limited = counts

            if requests <= 0 or fans <= 0:
                continue

            # `is not None` is a no-op (WHERE filters it) — for the type-checker.
            start = min(row.started_at for row in group)
            finishes = [row.finished_at for row in group if row.finished_at is not None]

            # Full-drain span (max finish) — the concurrency-queue tail p90 dropped.
            span_seconds = (max(finishes) - start).total_seconds()

            if span_seconds <= 0:
                continue

            serial_seconds = sum(
                (row.finished_at - row.started_at).total_seconds()
                for row in group
                if row.finished_at is not None
            )

            rps_list.append(requests / span_seconds)
            rpf_list.append(requests / fans)
            duration_list.append(span_seconds)
            time_per_fan_list.append(serial_seconds / fans)
            latency_list.append(serial_seconds / requests)
            throttle_list.append(rate_limited / requests)
            # Fan drain rate — the sizing signal (single ratio, self-corrects on
            # overrun: longer span → fewer fans next dispatch).
            fans_per_s_list.append(fans / span_seconds)

        if not rps_list:
            return None

        # Most recent fanout that has data (recent_ids is newest-first) — the
        # no-growth cap reference for the overrun damper.
        last_fans = 0
        for fanout_id in recent_ids:
            counts = counts_by_fanout.get(fanout_id)
            if counts:
                last_fans = counts[1]
                break

        # 5. Return medians
        return FanoutSignals(
            rps=median(rps_list),
            reqs_per_fan=median(rpf_list),
            duration_s=median(duration_list),
            latency_s=median(latency_list),
            time_per_fan_s=median(time_per_fan_list),
            throttle_rate=median(throttle_list),
            fans_per_s=median(fans_per_s_list),
            last_fans=last_fans,
        )

    def mark_stale_one(self, collect_task_id: uuid.UUID) -> bool:
        """Mark a single collect task as stale (DLQ path — SQS retries exhausted).

        No-ops if the task is already done (Lambda succeeded on a prior attempt).
        Returns True if the status was changed.
        """
        now = timezone.now()
        result = db.session.execute(
            update(CollectTask)
            .where(
                CollectTask.id == collect_task_id,
                CollectTask.status != TaskStatus.done,
            )
            .values(
                status=TaskStatus.stale,
                finished_at=now,
                finished_reason=TaskFinishedReason.timed_out,
            )
        )
        return cast(CursorResult, result).rowcount == 1

    def update_outcome(
        self,
        collect_task_id: uuid.UUID,
        *,
        fans_processed: int,
        fans_errors: int,
        fans_stale_tokens: int,
        fans_skipped: int,
        requests: int,
        requests_rate_limited: int,
        requests_skipped: int,
    ) -> None:
        """Persist collection outcome counters and mark the task as done."""
        db.session.execute(
            update(CollectTask)
            .where(CollectTask.id == collect_task_id)
            .values(
                status=TaskStatus.done,
                finished_at=timezone.now(),
                finished_reason=TaskFinishedReason.done,
                fans_processed=fans_processed,
                fans_errors=fans_errors,
                fans_stale_tokens=fans_stale_tokens,
                requests_rate_limited=requests_rate_limited,
                fans_skipped=fans_skipped,
                requests_skipped=requests_skipped,
                requests=requests,
            )
        )


class CollectTask(Model, kw_only=True):
    __tablename__ = "task_collect"
    __table_args__ = (
        Index(
            "ix_task_collect_dsp_client_name_started_at",
            "dsp_client_name",
            "started_at",
        ),
        Index(
            "ix_task_collect_fanout_task_id_dsp_client_name",
            "fanout_task_id",
            "dsp_client_name",
        ),
        Index("ix_task_collect_started_at_status", "started_at", "status"),
    )

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), primary_key=True, default_factory=uuid.uuid4, init=False
    )
    fanout_task_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True))
    dsp_client_name: Mapped[DSPClientName]
    started_at: Mapped[datetime]
    status: Mapped[TaskStatus]
    finished_at: Mapped[datetime | None] = mapped_column(default=None)
    finished_reason: Mapped[TaskFinishedReason | None] = mapped_column(default=None)
    fans_total: Mapped[int]
    fans_skipped: Mapped[int] = mapped_column(default=0)
    fans_processed: Mapped[int] = mapped_column(default=0)
    fans_errors: Mapped[int] = mapped_column(default=0)
    fans_stale_tokens: Mapped[int] = mapped_column(default=0)
    requests: Mapped[int] = mapped_column(default=0)
    requests_skipped: Mapped[int] = mapped_column(default=0)
    requests_rate_limited: Mapped[int] = mapped_column(default=0)

    query = CollectTaskQuery.as_descriptor()