The Grand Finale.

# budget/distributed_budget.py

"""
Hierarchical budget management for real-time bidding.

Provides strong consistency guarantees while minimizing
latency impact on the critical bidding path.
"""

from dataclasses import dataclass
from typing import Dict, Optional, Tuple
from datetime import datetime, timedelta, timezone
import asyncio
import logging
import time

logger = logging.getLogger(__name__)


@dataclass
class BudgetLease:
    """A local budget allocation."""
    advertiser_id: str
    campaign_id: str
    allocated_amount: float
    spent_amount: float
    expires_at: float  # deadline on this pod's monotonic clock (time.monotonic())
    
    @property
    def remaining(self) -> float:
        return self.allocated_amount - self.spent_amount
    
    @property
    def is_expired(self) -> bool:
        # Monotonic clock: wall-clock jumps (NTP steps) can't extend or cut a lease
        return time.monotonic() >= self.expires_at
    
    @property
    def utilization(self) -> float:
        if self.allocated_amount == 0:
            return 1.0
        return self.spent_amount / self.allocated_amount


class LocalBudgetManager:
    """
    Manages budget leases for a single server pod.
    
    Design principles:
    - Spend decisions are LOCAL (no network in hot path)
    - Background refresh keeps leases topped up
    - Conservative: rather reject bids than overspend
    """
    
    def __init__(
        self,
        region_allocator,
        pod_id: str,
        refresh_threshold: float = 0.7,  # Refresh at 70% utilization
        lease_duration: timedelta = timedelta(seconds=25)  # holder stops at 25s; the region reclaims after 30s
    ):
        self.region = region_allocator
        self.pod_id = pod_id
        self.refresh_threshold = refresh_threshold
        self.lease_duration = lease_duration
        
        self.leases: Dict[str, BudgetLease] = {}
        self._refresh_lock = asyncio.Lock()
        self._background: set = set()  # strong refs so fire-and-forget tasks aren't garbage-collected
    
    def _spawn(self, coro) -> None:
        """Fire-and-forget, but keep a reference until the task finishes."""
        task = asyncio.create_task(coro)
        self._background.add(task)
        task.add_done_callback(self._background.discard)
    
    async def can_spend(
        self,
        advertiser_id: str,
        campaign_id: str,
        amount: float
    ) -> Tuple[bool, str]:
        """
        Check if we can spend amount. MUST BE FAST (<1ms).
        
        Returns (can_spend, reason)
        """
        key = f"{advertiser_id}:{campaign_id}"
        lease = self.leases.get(key)
        
        # No lease? Try to get one (but don't block hot path)
        if lease is None:
            # Trigger background refresh
            self._spawn(self._refresh_lease(advertiser_id, campaign_id))
            return False, "no_budget_lease"
        
        # Lease expired?
        if lease.is_expired:
            self._spawn(self._refresh_lease(advertiser_id, campaign_id))
            return False, "lease_expired"
        
        # Not enough remaining?
        if lease.remaining < amount:
            self._spawn(self._refresh_lease(advertiser_id, campaign_id))
            return False, "insufficient_budget"
        
        # Check if we need background refresh
        if lease.utilization > self.refresh_threshold:
            self._spawn(self._refresh_lease(advertiser_id, campaign_id))
        
        return True, "approved"
    
    async def record_spend(
        self,
        advertiser_id: str,
        campaign_id: str,
        amount: float
    ) -> bool:
        """
        Record spend against local lease.
        Call this AFTER winning auction, BEFORE responding.
        """
        key = f"{advertiser_id}:{campaign_id}"
        lease = self.leases.get(key)
        
        if lease is None or lease.remaining < amount:
            # This shouldn't happen if can_spend was called
            logger.warning(f"Spend rejected: {key}, amount={amount}")
            return False
        
        lease.spent_amount += amount
        
        # Async report to region (for reconciliation)
        self._spawn(
            self.region.report_spend(advertiser_id, campaign_id, amount)
        )
        
        return True
    
    async def _refresh_lease(
        self,
        advertiser_id: str,
        campaign_id: str
    ):
        """Background refresh of budget lease."""
        key = f"{advertiser_id}:{campaign_id}"
        
        async with self._refresh_lock:
            try:
                # Request new allocation from region
                allocation = await self.region.request_allocation(
                    advertiser_id=advertiser_id,
                    campaign_id=campaign_id,
                    pod_id=self.pod_id,
                    requested_amount=self._calculate_request_amount(key)
                )
                
                if allocation.granted_amount > 0:
                    self.leases[key] = BudgetLease(
                        advertiser_id=advertiser_id,
                        campaign_id=campaign_id,
                        allocated_amount=allocation.granted_amount,
                        spent_amount=0,
                        expires_at=time.monotonic() + self.lease_duration.total_seconds()
                    )
                    
            except Exception as e:
                logger.error(f"Failed to refresh lease {key}: {e}")


class RegionBudgetAllocator:
    """
    Manages budget allocation for a region.
    
    Coordinates between:
    - Global coordinator (gets region's share)
    - Local pods (distributes to pods)
    """
    
    def __init__(
        self,
        global_coordinator,
        region_id: str,
        redis_cluster
    ):
        self.global_coord = global_coordinator
        self.region_id = region_id
        self.redis = redis_cluster
    
    async def request_allocation(
        self,
        advertiser_id: str,
        campaign_id: str,
        pod_id: str,
        requested_amount: float
    ) -> 'AllocationResult':
        """
        Allocate budget to a pod.
        Uses Redis for atomic operations within region.
        """
        key = f"budget:{advertiser_id}:{campaign_id}:{self.region_id}"
        
        # Lua script for atomic check-and-decrement
        lua_script = """
        local available = tonumber(redis.call('GET', KEYS[1]) or '0')
        local requested = tonumber(ARGV[1])
        
        if available <= 0 then
            return 0
        end
        
        local granted = math.min(available, requested)
        -- INCRBYFLOAT, not DECRBY: amounts are fractional dollars
        -- (or store integer cents and use DECRBY)
        redis.call('INCRBYFLOAT', KEYS[1], -granted)
        return tostring(granted)  -- Lua numbers become integers in replies
        """
        
        granted = await self.redis.eval(
            lua_script,
            1, key, requested_amount
        )
        
        granted = float(granted)
        
        if granted == 0:
            # Try to get more from global
            await self._request_global_refill(advertiser_id, campaign_id)
            
            # Retry once
            granted = float(await self.redis.eval(
                lua_script,
                1, key, requested_amount
            ))
        
        return AllocationResult(
            granted_amount=granted,
            advertiser_id=advertiser_id,
            campaign_id=campaign_id
        )

# bidding/dsp_gateway.py

"""
DSP Gateway: Manages communication with demand-side platforms.

Challenges:
- Must query 5-10 DSPs in parallel
- Each DSP has different latency characteristics
- Total budget: 70ms (must leave margin for auction)
- Partial responses are OK (bid with what we have)
"""

from dataclasses import dataclass, field
from typing import List, Dict, Optional, Set
from datetime import datetime, timezone
import asyncio
import logging

logger = logging.getLogger(__name__)


@dataclass
class DSPConfig:
    """Configuration for a DSP."""
    dsp_id: str
    endpoint: str
    timeout_ms: int = 70
    
    # Circuit breaker settings
    failure_threshold: int = 5      # Failures before opening
    success_threshold: int = 3      # Successes before closing
    half_open_timeout_s: int = 30   # Time before retrying
    
    # Performance tracking
    p99_latency_ms: float = 50
    error_rate: float = 0.01


@dataclass
class DSPBid:
    """A bid from a DSP."""
    dsp_id: str
    bid_id: str
    price: float
    ad_markup: str
    latency_ms: float
    received_at: datetime


@dataclass
class FanOutResult:
    """Results from DSP fan-out."""
    bids: List[DSPBid]
    timeouts: List[str]          # DSP IDs that timed out
    errors: List[str]            # DSP IDs that errored
    circuit_open: List[str]      # DSP IDs with open circuits
    total_latency_ms: float


class CircuitBreaker:
    """
    Circuit breaker for a single DSP.
    
    States:
    - CLOSED: Normal operation, requests flow through
    - OPEN: DSP is unhealthy, fail fast (don't even try)
    - HALF_OPEN: Testing if DSP recovered
    """
    
    def __init__(self, config: DSPConfig):
        self.config = config
        self.state = "CLOSED"
        self.failure_count = 0
        self.success_count = 0
        self.last_failure_time: Optional[datetime] = None
    
    def can_execute(self) -> bool:
        if self.state == "CLOSED":
            return True
        
        if self.state == "OPEN":
            # Check if we should try half-open
            if self.last_failure_time:
                elapsed = (datetime.now(timezone.utc) - self.last_failure_time).total_seconds()
                if elapsed > self.config.half_open_timeout_s:
                    self.state = "HALF_OPEN"
                    return True
            return False
        
        if self.state == "HALF_OPEN":
            return True
        
        return False
    
    def record_success(self):
        if self.state == "HALF_OPEN":
            self.success_count += 1
            if self.success_count >= self.config.success_threshold:
                self.state = "CLOSED"
                self.failure_count = 0
                self.success_count = 0
        else:
            self.failure_count = 0
    
    def record_failure(self):
        self.failure_count += 1
        self.last_failure_time = datetime.now(timezone.utc)
        
        if self.state == "HALF_OPEN":
            self.state = "OPEN"
            self.success_count = 0
        elif self.failure_count >= self.config.failure_threshold:
            self.state = "OPEN"


class DSPGateway:
    """
    Gateway for querying DSPs with timeout and failure handling.
    """
    
    def __init__(
        self,
        http_client,
        dsp_configs: List[DSPConfig],
        global_timeout_ms: int = 70
    ):
        self.http = http_client
        self.dsps = {cfg.dsp_id: cfg for cfg in dsp_configs}
        self.global_timeout_ms = global_timeout_ms
        
        # Circuit breakers per DSP
        self.breakers = {
            cfg.dsp_id: CircuitBreaker(cfg) 
            for cfg in dsp_configs
        }
        
        # Metrics
        self.metrics = DSPMetrics()
    
    async def fan_out(
        self,
        request: 'BidRequest',
        target_dsps: List[str]
    ) -> FanOutResult:
        """
        Query multiple DSPs in parallel with timeout.
        
        Key design decisions:
        1. Start all requests immediately (parallel)
        2. Use global timeout (not per-DSP)
        3. Collect responses as they arrive
        4. Return whatever we have when timeout hits
        """
        start_time = datetime.now(timezone.utc)
        
        result = FanOutResult(
            bids=[],
            timeouts=[],
            errors=[],
            circuit_open=[],
            total_latency_ms=0
        )
        
        # Filter out DSPs with open circuits
        active_dsps = []
        for dsp_id in target_dsps:
            breaker = self.breakers.get(dsp_id)
            if breaker and breaker.can_execute():
                active_dsps.append(dsp_id)
            else:
                result.circuit_open.append(dsp_id)
        
        if not active_dsps:
            return result
        
        # Create tasks for all DSPs
        tasks = {
            dsp_id: asyncio.create_task(
                self._query_dsp(dsp_id, request)
            )
            for dsp_id in active_dsps
        }
        
        # Wait with timeout, collecting results as they arrive
        pending = set(tasks.values())
        deadline = self.global_timeout_ms / 1000
        
        try:
            while pending:
                # Calculate remaining time
                elapsed = (datetime.now(timezone.utc) - start_time).total_seconds()
                remaining = deadline - elapsed
                
                if remaining <= 0:
                    # Timeout! Cancel remaining and break
                    for task in pending:
                        task.cancel()
                    break
                
                # Wait for next completion or timeout
                done, pending = await asyncio.wait(
                    pending,
                    timeout=remaining,
                    return_when=asyncio.FIRST_COMPLETED
                )
                
                # Process completed tasks
                for task in done:
                    dsp_id = self._get_dsp_id_for_task(task, tasks)
                    
                    try:
                        bid = task.result()
                        if bid:
                            result.bids.append(bid)
                            self.breakers[dsp_id].record_success()
                    except asyncio.TimeoutError:
                        result.timeouts.append(dsp_id)
                        self.breakers[dsp_id].record_failure()
                    except Exception as e:
                        result.errors.append(dsp_id)
                        self.breakers[dsp_id].record_failure()
                        logger.warning(f"DSP {dsp_id} error: {e}")
        
        finally:
            # Cancel any still pending: they missed the deadline, which
            # counts as a failure for that DSP's circuit breaker
            for task in pending:
                task.cancel()
                dsp_id = self._get_dsp_id_for_task(task, tasks)
                result.timeouts.append(dsp_id)
                self.breakers[dsp_id].record_failure()
        
        result.total_latency_ms = (
            datetime.now(timezone.utc) - start_time
        ).total_seconds() * 1000
        
        # Log metrics
        self.metrics.record_fanout(result)
        
        return result
    
    async def _query_dsp(
        self,
        dsp_id: str,
        request: 'BidRequest'
    ) -> Optional[DSPBid]:
        """Query a single DSP."""
        config = self.dsps[dsp_id]
        start = datetime.now(timezone.utc)
        
        try:
            response = await self.http.post(
                config.endpoint,
                json=request.to_dsp_format(),
                timeout=config.timeout_ms / 1000
            )
            
            latency = (datetime.now(timezone.utc) - start).total_seconds() * 1000
            
            if response.status == 204:
                # No bid
                return None
            
            bid_data = response.json()
            
            return DSPBid(
                dsp_id=dsp_id,
                bid_id=bid_data['id'],
                price=bid_data['price'],
                ad_markup=bid_data['adm'],
                latency_ms=latency,
                received_at=datetime.now(timezone.utc)
            )
            
        except asyncio.TimeoutError:
            raise
        except Exception as e:
            logger.error(f"DSP {dsp_id} request failed: {e}")
            raise

# audit/event_pipeline.py

"""
Audit event schema and processing.

Every bid event follows this schema for consistency
across billing, analytics, and compliance.
"""

from dataclasses import dataclass, asdict
from typing import List, Dict, Optional, Any
from datetime import datetime
import json


@dataclass
class BidEvent:
    """
    Immutable record of a bid transaction.
    
    This is the source of truth for:
    - Billing reconciliation
    - Analytics dashboards  
    - Fraud investigation
    - Compliance audits
    """
    # Identity
    event_id: str              # UUID, globally unique
    event_type: str            # "bid_request", "bid_response", "auction_result"
    event_time: datetime       # When this happened (event time, not processing time)
    
    # Request context
    request_id: str            # Correlates all events for one request
    exchange_id: str           # Which ad exchange
    publisher_id: str          # Which publisher
    
    # Advertiser context
    advertiser_id: str         # For partitioning and billing
    campaign_id: str
    
    # Auction details
    bid_price: Optional[float] = None
    win_price: Optional[float] = None  # Second price (what they actually pay)
    is_winner: bool = False
    
    # DSP details
    dsp_id: Optional[str] = None
    dsp_latency_ms: Optional[float] = None
    
    # User context (privacy-compliant)
    user_id_hash: Optional[str] = None  # Hashed, not raw
    geo_country: Optional[str] = None
    device_type: Optional[str] = None
    
    # Technical metadata
    region: str = ""
    pod_id: str = ""
    processing_time_ms: float = 0
    
    def to_json(self) -> str:
        """Serialize for Kafka."""
        data = asdict(self)
        data['event_time'] = self.event_time.isoformat()
        return json.dumps(data)
    
    @property
    def partition_key(self) -> str:
        """Kafka partition key ensures ordering per advertiser."""
        return self.advertiser_id


# Kafka topic configuration
KAFKA_TOPIC_CONFIG = {
    'name': 'bid-events',
    'partitions': 256,          # High partition count for parallelism
    'replication_factor': 3,    # Durability
    'retention_ms': 7 * 24 * 60 * 60 * 1000,  # 7 days hot
    'cleanup_policy': 'delete',
    
    # Compression settings for efficiency
    'compression_type': 'zstd',  # Best compression ratio
    'max_message_bytes': 1048576,  # 1MB max
}

# fraud/realtime_detector.py

"""
Real-time fraud detection in the bidding pipeline.

Constraint: Must complete in <5ms including all tiers.
"""

from dataclasses import dataclass
from typing import Tuple
import numpy as np


@dataclass  
class FraudSignals:
    """Features extracted for fraud detection."""
    ip_address: str
    user_agent: str
    device_id: str
    geo_country: str
    request_rate_1m: int    # Requests from this IP in last minute
    request_rate_1h: int    # Requests from this IP in last hour
    user_age_days: int      # How long we've seen this user
    session_duration: float # Current session length
    click_rate: float       # Historical click-through rate


class RealTimeFraudDetector:
    """
    Tiered fraud detection for real-time bidding.
    """
    
    def __init__(
        self,
        blocklist_store,
        feature_store,
        ml_model,
        rate_limiter
    ):
        self.blocklist = blocklist_store  # Redis set
        self.features = feature_store      # Redis hash
        self.model = ml_model              # Pre-loaded XGBoost
        self.limiter = rate_limiter
    
    async def evaluate(
        self,
        signals: FraudSignals
    ) -> Tuple[str, float]:
        """
        Evaluate fraud risk.
        
        Returns: (decision, confidence)
        - decision: "block", "flag", "allow"
        - confidence: 0.0 to 1.0
        """
        
        # TIER 1: Rule-based checks (<0.5ms)
        tier1_result = await self._tier1_rules(signals)
        if tier1_result == "block":
            return "block", 0.99
        
        # TIER 2: ML scoring (<2ms)
        fraud_score = await self._tier2_ml_score(signals)
        
        if fraud_score > 0.95:
            return "block", fraud_score
        elif fraud_score > 0.5:
            return "flag", fraud_score
        else:
            return "allow", 1 - fraud_score
    
    async def _tier1_rules(self, signals: FraudSignals) -> str:
        """
        Fast rule-based checks.
        
        Rules are simple, but catch obvious fraud instantly.
        """
        # Check IP blocklist
        if await self.blocklist.contains(signals.ip_address):
            return "block"
        
        # Check rate limits
        if signals.request_rate_1m > 100:  # >100 requests/minute from one IP
            return "block"
        
        # Check for data center IP ranges
        if self._is_datacenter_ip(signals.ip_address):
            return "block"
        
        return "allow"
    
    async def _tier2_ml_score(self, signals: FraudSignals) -> float:
        """
        ML-based fraud scoring.
        
        Uses pre-computed features and lightweight model.
        """
        # Get historical features from Redis
        historical = await self.features.get(signals.device_id)
        
        # Build feature vector
        features = np.array([
            signals.request_rate_1m,
            signals.request_rate_1h,
            signals.user_age_days,
            signals.session_duration,
            signals.click_rate,
            historical.get('conversion_rate', 0),
            historical.get('avg_session_length', 0),
            historical.get('fraud_score_history', 0),
        ])
        
        # Score with pre-loaded model (XGBoost inference is fast)
        score = self.model.predict_proba(features.reshape(1, -1))[0][1]
        
        return score