"""Polymarket liquidity rewards tracker for hz.run().

Tracks maker fills, estimates rebates at a configurable rate, and
provides epoch-based reward accumulation.

Usage:
    hz.run(
        pipeline=[model, quoter, hz.rewards_tracker()],
        ...
    )
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Any, Callable

from horizon.context import Context


@dataclass
class EpochRewards:
    """Per-market reward accumulator for one epoch."""

    market_id: str = ""
    fills_scored: int = 0
    estimated_usdc: float = 0.0
    total_notional: float = 0.0

    def as_dict(self) -> dict[str, Any]:
        return {
            "market_id": self.market_id,
            "fills_scored": self.fills_scored,
            "estimated_usdc": self.estimated_usdc,
            "total_notional": self.total_notional,
        }


class RewardsTracker:
    """Tracks maker fills and estimates liquidity rebates.

    Args:
        rebate_rate: Rebate rate as fraction of notional (default 0.0002 = 2bps).
        check_scoring: If True, calls engine.check_order_scoring() per fill.
    """

    def __init__(
        self,
        rebate_rate: float = 0.0002,
        check_scoring: bool = False,
    ) -> None:
        self.rebate_rate = rebate_rate
        self.check_scoring = check_scoring
        self._epochs: dict[str, EpochRewards] = {}
        self._seen_fills: set[str] = set()
        self._total_usdc: float = 0.0
        self._epoch_start: float = time.time()

    def record_fill(
        self,
        market_id: str,
        fill_id: str,
        notional: float,
        is_maker: bool = True,
        scoring: bool = True,
    ) -> float:
        """Record a fill and return estimated rebate.

        Args:
            market_id: Market identifier.
            fill_id: Unique fill ID for deduplication.
            notional: Fill notional value (price * size).
            is_maker: Whether this was a maker fill.
            scoring: Whether the order qualified for rewards.

        Returns:
            Estimated rebate in USDC for this fill (0 if not eligible).
        """
        if fill_id in self._seen_fills:
            return 0.0

        self._seen_fills.add(fill_id)

        if not is_maker or not scoring or notional <= 0:
            return 0.0

        rebate = notional * self.rebate_rate

        if market_id not in self._epochs:
            self._epochs[market_id] = EpochRewards(market_id=market_id)
        epoch = self._epochs[market_id]
        epoch.fills_scored += 1
        epoch.estimated_usdc += rebate
        epoch.total_notional += notional

        self._total_usdc += rebate
        return rebate

    @property
    def epoch_rewards(self) -> dict[str, EpochRewards]:
        return dict(self._epochs)

    @property
    def total_usdc(self) -> float:
        return self._total_usdc

    def reset_epoch(self) -> dict[str, EpochRewards]:
        """Reset and return the current epoch rewards."""
        old = dict(self._epochs)
        self._epochs.clear()
        self._epoch_start = time.time()
        return old


def rewards_tracker(
    rebate_rate: float = 0.0002,
    check_scoring: bool = False,
) -> Callable[[Context], dict[str, Any]]:
    """Create a rewards tracker pipeline function.

    Monitors fills from the engine, estimates rebates, and injects
    reward metrics into the pipeline output.

    Args:
        rebate_rate: Rebate rate as fraction of notional (default 2bps).
        check_scoring: If True, check order scoring via Polymarket API.

    Returns:
        Pipeline function: (Context) -> dict with keys:
            epoch_rewards, total_usdc, reward_markets_count
    """
    from horizon._horizon import auth_require_pro

    auth_require_pro()

    tracker = RewardsTracker(rebate_rate=rebate_rate, check_scoring=check_scoring)
    _last_fill_count: dict[str, int] = {}

    def _rewards_tracker(ctx: Context) -> dict[str, Any]:
        # Check for new fills in the engine
        engine = ctx.engine
        if engine is not None:
            fills = engine.fills()
            prev_count = _last_fill_count.get("__total__", 0)
            if len(fills) > prev_count:
                new_fills = fills[prev_count:]
                _last_fill_count["__total__"] = len(fills)

                for fill in new_fills:
                    notional = fill.price * fill.size
                    is_maker = True  # Assume maker for MM fills

                    scoring = True
                    if check_scoring and engine is not None:
                        try:
                            scoring = engine.check_order_scoring(fill.order_id)
                        except Exception:
                            scoring = True  # Assume scoring on API error

                    tracker.record_fill(
                        market_id=fill.market_id,
                        fill_id=fill.fill_id,
                        notional=notional,
                        is_maker=is_maker,
                        scoring=scoring,
                    )

        return {
            "epoch_rewards": {
                k: v.as_dict() for k, v in tracker.epoch_rewards.items()
            },
            "total_usdc": tracker.total_usdc,
            "reward_markets_count": len(tracker.epoch_rewards),
        }

    _rewards_tracker.__name__ = "rewards_tracker"
    return _rewards_tracker
