"""Abstract store interface and in-memory implementation for referral data."""

from __future__ import annotations

import abc
from datetime import datetime, timezone
from uuid import UUID

from referral.models import (
    Referral,
    ReferralCode,
    ReferralReward,
    ReferralStats,
    ReferralStatus,
    RewardStatus,
)


class ReferralStore(abc.ABC):
    """Abstract base class defining the referral persistence contract.

    Production implementations should use asyncpg; the in-memory variant
    is provided for testing and prototyping.
    """

    # -- codes --

    @abc.abstractmethod
    async def save_code(self, code: ReferralCode) -> ReferralCode: ...

    @abc.abstractmethod
    async def get_code(self, code: str) -> ReferralCode | None: ...

    @abc.abstractmethod
    async def get_code_by_id(self, code_id: UUID) -> ReferralCode | None: ...

    @abc.abstractmethod
    async def get_user_codes(self, user_id: str) -> list[ReferralCode]: ...

    @abc.abstractmethod
    async def update_code(self, code: ReferralCode) -> ReferralCode: ...

    # -- referrals --

    @abc.abstractmethod
    async def save_referral(self, referral: Referral) -> Referral: ...

    @abc.abstractmethod
    async def get_referral(self, referral_id: UUID) -> Referral | None: ...

    @abc.abstractmethod
    async def list_referrals(
        self,
        *,
        referrer_id: str | None = None,
        referred_id: str | None = None,
        status: ReferralStatus | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> list[Referral]: ...

    @abc.abstractmethod
    async def update_referral(self, referral: Referral) -> Referral: ...

    @abc.abstractmethod
    async def count_user_referrals(self, user_id: str) -> int: ...

    # -- rewards --

    @abc.abstractmethod
    async def save_reward(self, reward: ReferralReward) -> ReferralReward: ...

    @abc.abstractmethod
    async def get_user_rewards(self, user_id: str) -> list[ReferralReward]: ...

    @abc.abstractmethod
    async def get_pending_rewards(self) -> list[ReferralReward]: ...

    @abc.abstractmethod
    async def update_reward(self, reward: ReferralReward) -> ReferralReward: ...

    # -- stats --

    @abc.abstractmethod
    async def get_stats(self, user_id: str) -> ReferralStats: ...

    @abc.abstractmethod
    async def get_leaderboard(self, limit: int = 10) -> list[ReferralStats]: ...

    # -- maintenance --

    @abc.abstractmethod
    async def expire_stale(self, before: datetime) -> int: ...


class InMemoryReferralStore(ReferralStore):
    """In-memory referral store for tests and local development."""

    def __init__(self) -> None:
        self.codes: dict[str, ReferralCode] = {}
        self.codes_by_id: dict[UUID, ReferralCode] = {}
        self.referrals: dict[UUID, Referral] = {}
        self.rewards: dict[UUID, ReferralReward] = {}

    # -- codes --

    async def save_code(self, code: ReferralCode) -> ReferralCode:
        self.codes[code.code] = code
        self.codes_by_id[code.id] = code
        return code

    async def get_code(self, code: str) -> ReferralCode | None:
        return self.codes.get(code)

    async def get_code_by_id(self, code_id: UUID) -> ReferralCode | None:
        return self.codes_by_id.get(code_id)

    async def get_user_codes(self, user_id: str) -> list[ReferralCode]:
        return [c for c in self.codes.values() if c.user_id == user_id]

    async def update_code(self, code: ReferralCode) -> ReferralCode:
        self.codes[code.code] = code
        self.codes_by_id[code.id] = code
        return code

    # -- referrals --

    async def save_referral(self, referral: Referral) -> Referral:
        self.referrals[referral.id] = referral
        return referral

    async def get_referral(self, referral_id: UUID) -> Referral | None:
        return self.referrals.get(referral_id)

    async def list_referrals(
        self,
        *,
        referrer_id: str | None = None,
        referred_id: str | None = None,
        status: ReferralStatus | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> list[Referral]:
        results = list(self.referrals.values())
        if referrer_id is not None:
            results = [r for r in results if r.referrer_id == referrer_id]
        if referred_id is not None:
            results = [r for r in results if r.referred_id == referred_id]
        if status is not None:
            results = [r for r in results if r.status == status]
        results.sort(key=lambda r: r.created_at, reverse=True)
        return results[offset : offset + limit]

    async def update_referral(self, referral: Referral) -> Referral:
        self.referrals[referral.id] = referral
        return referral

    async def count_user_referrals(self, user_id: str) -> int:
        return sum(1 for r in self.referrals.values() if r.referrer_id == user_id)

    # -- rewards --

    async def save_reward(self, reward: ReferralReward) -> ReferralReward:
        self.rewards[reward.id] = reward
        return reward

    async def get_user_rewards(self, user_id: str) -> list[ReferralReward]:
        return [r for r in self.rewards.values() if r.user_id == user_id]

    async def get_pending_rewards(self) -> list[ReferralReward]:
        return [r for r in self.rewards.values() if r.status == RewardStatus.PENDING]

    async def update_reward(self, reward: ReferralReward) -> ReferralReward:
        self.rewards[reward.id] = reward
        return reward

    # -- stats --

    async def get_stats(self, user_id: str) -> ReferralStats:
        user_referrals = [r for r in self.referrals.values() if r.referrer_id == user_id]
        user_rewards = [r for r in self.rewards.values() if r.user_id == user_id]
        return ReferralStats(
            user_id=user_id,
            total_referrals=len(user_referrals),
            completed=sum(
                1
                for r in user_referrals
                if r.status in {ReferralStatus.COMPLETED, ReferralStatus.REWARDED}
            ),
            pending=sum(1 for r in user_referrals if r.status == ReferralStatus.PENDING),
            expired=sum(1 for r in user_referrals if r.status == ReferralStatus.EXPIRED),
            total_rewards_earned=sum(
                r.amount for r in user_rewards if r.status == RewardStatus.GRANTED
            ),
        )

    async def get_leaderboard(self, limit: int = 10) -> list[ReferralStats]:
        user_ids = {r.referrer_id for r in self.referrals.values()}
        stats = [await self.get_stats(uid) for uid in user_ids]
        stats.sort(key=lambda s: s.completed, reverse=True)
        return stats[:limit]

    # -- maintenance --

    async def expire_stale(self, before: datetime) -> int:
        count = 0
        for referral in self.referrals.values():
            if (
                referral.status == ReferralStatus.PENDING
                and referral.created_at < before
            ):
                referral.status = ReferralStatus.EXPIRED
                count += 1
        return count
