Late-Arriving Data and Correctness.

# late_data/watermark_strategies.py

"""
Watermark generation strategies.

Different strategies for different use cases.
"""

from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Optional, Dict
from abc import ABC, abstractmethod


class WatermarkStrategy(ABC):
    """Base class for watermark strategies."""
    
    @abstractmethod
    def on_event(self, event_time: datetime) -> datetime:
        """Update watermark based on new event."""
        pass
    
    @abstractmethod
    def get_watermark(self) -> datetime:
        """Get current watermark."""
        pass


class BoundedOutOfOrdernessWatermark(WatermarkStrategy):
    """
    Most common strategy: Allow fixed lateness.
    
    Watermark = max_event_time - max_lateness
    
    Use when: Lateness is bounded and predictable
    Example: Mobile app with max 1 hour offline
    """
    
    def __init__(self, max_lateness: timedelta):
        self.max_lateness = max_lateness
        self.max_event_time: Optional[datetime] = None
    
    def on_event(self, event_time: datetime) -> datetime:
        if self.max_event_time is None:
            self.max_event_time = event_time
        else:
            self.max_event_time = max(self.max_event_time, event_time)
        
        return self.get_watermark()
    
    def get_watermark(self) -> datetime:
        if self.max_event_time is None:
            return datetime.min.replace(tzinfo=timezone.utc)
        return self.max_event_time - self.max_lateness


class PerKeyWatermark(WatermarkStrategy):
    """
    Track watermark per key (e.g., per user, per device).
    
    Global watermark = min(all key watermarks)
    
    Use when: Different sources have different lateness
    Example: Some users always online, some often offline
    """
    
    def __init__(self, max_lateness: timedelta, idle_timeout: timedelta):
        self.max_lateness = max_lateness
        self.idle_timeout = idle_timeout
        self.key_watermarks: Dict[str, datetime] = {}
        self.key_last_seen: Dict[str, datetime] = {}
    
    def on_event_for_key(
        self,
        key: str,
        event_time: datetime,
        processing_time: datetime
    ) -> datetime:
        # Update key's max event time
        if key not in self.key_watermarks:
            self.key_watermarks[key] = event_time
        else:
            self.key_watermarks[key] = max(
                self.key_watermarks[key],
                event_time
            )
        
        self.key_last_seen[key] = processing_time
        
        return self.get_watermark()
    
    def get_watermark(self) -> datetime:
        now = datetime.now(timezone.utc)
        
        # Filter out idle keys (haven't seen events recently)
        active_watermarks = []
        for key, wm in self.key_watermarks.items():
            last_seen = self.key_last_seen.get(key, now)
            if now - last_seen < self.idle_timeout:
                active_watermarks.append(wm - self.max_lateness)
        
        if not active_watermarks:
            return datetime.min.replace(tzinfo=timezone.utc)
        
        # Global watermark is minimum of active key watermarks
        return min(active_watermarks)


class PunctuatedWatermark(WatermarkStrategy):
    """
    Watermarks embedded in the data stream.
    
    Use when: Source can guarantee ordering
    Example: Database CDC with transaction IDs
    """
    
    def __init__(self):
        self.current_watermark = datetime.min.replace(tzinfo=timezone.utc)
    
    def on_event(self, event_time: datetime) -> datetime:
        # Regular events don't advance watermark
        return self.current_watermark
    
    def on_watermark_event(self, watermark_time: datetime) -> datetime:
        # Special watermark events advance the watermark
        self.current_watermark = max(
            self.current_watermark,
            watermark_time
        )
        return self.current_watermark
    
    def get_watermark(self) -> datetime:
        return self.current_watermark


class ProcessingTimeWatermark(WatermarkStrategy):
    """
    Use processing time as watermark (ignore event time).
    
    Use when: Event time doesn't matter or isn't reliable
    Example: Log processing where arrival order is fine
    """
    
    def __init__(self, delay: timedelta = timedelta(seconds=0)):
        self.delay = delay
    
    def on_event(self, event_time: datetime) -> datetime:
        return self.get_watermark()
    
    def get_watermark(self) -> datetime:
        return datetime.now(timezone.utc) - self.delay

# late_data/handling_strategies.py

"""
Strategies for handling late-arriving data.
"""

from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from typing import Dict, Any, List, Optional, Callable
from enum import Enum
import logging

logger = logging.getLogger(__name__)


class LateDataAction(Enum):
    """What to do with late data."""
    DROP = "drop"                    # Ignore late events
    SIDE_OUTPUT = "side_output"      # Route to separate stream
    UPDATE_RESULT = "update_result"  # Update the window result
    RECOMPUTE = "recompute"          # Trigger full recomputation


@dataclass
class LateDataConfig:
    """Configuration for late data handling."""
    allowed_lateness: timedelta
    action: LateDataAction
    side_output_topic: Optional[str] = None
    max_updates_per_window: int = 10


@dataclass
class WindowResult:
    """Result from a window computation."""
    window_start: datetime
    window_end: datetime
    result: Dict[str, Any]
    is_final: bool
    update_count: int = 0
    last_updated: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


class LateDataHandler:
    """
    Handles late-arriving events in windowed computations.
    """
    
    def __init__(self, config: LateDataConfig):
        self.config = config
        self.window_results: Dict[str, WindowResult] = {}
        self.late_events: List[Dict] = []
        
        # Metrics
        self.on_time_count = 0
        self.late_accepted_count = 0
        self.late_dropped_count = 0
    
    def process_event(
        self,
        event: Dict,
        watermark: datetime,
        aggregate_fn: Callable
    ) -> Optional[Dict]:
        """
        Process an event, handling lateness appropriately.
        """
        
        event_time = datetime.fromisoformat(event["event_time"])
        window_key = self._get_window_key(event_time)
        window_start, window_end = self._parse_window_key(window_key)
        
        # Late = its window has already fired (watermark passed window end).
        # The watermark already subtracts the out-of-orderness bound, so
        # don't compare event_time to it directly.
        is_late = window_end <= watermark
        
        if not is_late:
            # On-time event: normal processing
            self.on_time_count += 1
            return self._update_window(window_key, event, aggregate_fn, is_late=False)
        
        # Event is late: how far the watermark is past the window end
        lateness = watermark - window_end
        
        # Check if within allowed lateness
        if lateness <= self.config.allowed_lateness:
            return self._handle_late_within_tolerance(
                window_key, event, aggregate_fn, lateness
            )
        else:
            return self._handle_late_beyond_tolerance(event, lateness)
    
    def _handle_late_within_tolerance(
        self,
        window_key: str,
        event: Dict,
        aggregate_fn: Callable,
        lateness: timedelta
    ) -> Optional[Dict]:
        """Handle event that's late but within tolerance."""
        
        self.late_accepted_count += 1
        
        if self.config.action == LateDataAction.DROP:
            logger.debug(f"Dropping late event (within tolerance): {lateness}")
            return None
        
        elif self.config.action == LateDataAction.SIDE_OUTPUT:
            self.late_events.append({
                "event": event,
                "lateness_seconds": lateness.total_seconds(),
                "action": "side_output"
            })
            return None
        
        elif self.config.action == LateDataAction.UPDATE_RESULT:
            # Check update limit
            existing = self.window_results.get(window_key)
            if existing and existing.update_count >= self.config.max_updates_per_window:
                logger.warning(
                    f"Window {window_key} exceeded max updates, dropping event"
                )
                return None
            
            return self._update_window(window_key, event, aggregate_fn, is_late=True)
        
        return None
    
    def _handle_late_beyond_tolerance(
        self,
        event: Dict,
        lateness: timedelta
    ) -> Optional[Dict]:
        """Handle event that's later than allowed tolerance."""
        
        self.late_dropped_count += 1
        
        logger.info(
            f"Dropping very late event: {lateness} exceeds "
            f"{self.config.allowed_lateness}"
        )
        
        # Always side-output very late events for analysis
        self.late_events.append({
            "event": event,
            "lateness_seconds": lateness.total_seconds(),
            "action": "dropped_beyond_tolerance"
        })
        
        return None
    
    def _update_window(
        self,
        window_key: str,
        event: Dict,
        aggregate_fn: Callable,
        is_late: bool
    ) -> Dict:
        """Update or create window result."""
        
        window_start, window_end = self._parse_window_key(window_key)
        
        existing = self.window_results.get(window_key)
        
        if existing:
            # Update existing window
            new_result = aggregate_fn(existing.result, event)
            existing.result = new_result
            existing.update_count += 1
            existing.last_updated = datetime.now(timezone.utc)
            existing.is_final = False  # No longer final since it was updated
            
            return {
                "window_key": window_key,
                "result": new_result,
                "is_update": True,
                "is_late_update": is_late,
                "update_count": existing.update_count
            }
        else:
            # New window
            result = aggregate_fn(None, event)
            self.window_results[window_key] = WindowResult(
                window_start=window_start,
                window_end=window_end,
                result=result,
                is_final=False,
                update_count=0
            )
            
            return {
                "window_key": window_key,
                "result": result,
                "is_update": False,
                "is_late_update": is_late,
                "update_count": 0
            }
    
    def finalize_window(self, window_key: str) -> Optional[WindowResult]:
        """Mark a window as final (watermark passed window end)."""
        
        if window_key in self.window_results:
            result = self.window_results[window_key]
            result.is_final = True
            return result
        return None
    
    def _get_window_key(self, event_time: datetime) -> str:
        """Window key for an event time (1-hour tumbling), e.g. 2024-01-15-14."""
        return event_time.strftime("%Y-%m-%d-%H")
    
    def _parse_window_key(self, key: str) -> tuple:
        """Parse window key back to start/end times."""
        start = datetime.strptime(key, "%Y-%m-%d-%H").replace(tzinfo=timezone.utc)
        return start, start + timedelta(hours=1)
    
    def get_metrics(self) -> Dict:
        """Get late data handling metrics."""
        total = self.on_time_count + self.late_accepted_count + self.late_dropped_count
        
        return {
            "on_time_count": self.on_time_count,
            "late_accepted_count": self.late_accepted_count,
            "late_dropped_count": self.late_dropped_count,
            "on_time_rate": self.on_time_count / total if total > 0 else 0,
            "late_rate": (self.late_accepted_count + self.late_dropped_count) / total if total > 0 else 0
        }

# late_data/state_cleanup.py

"""
Managing state with allowed lateness.

Problem: If we keep window state forever to handle late data,
         memory grows unbounded.

Solution: Use allowed lateness + cleanup timers.
"""

from datetime import datetime, timedelta, timezone
from typing import Dict, Any, Optional
import logging

logger = logging.getLogger(__name__)


class WindowStateManager:
    """
    Manages window state with cleanup for late data handling.
    
    State lifecycle:
    1. Window created when first event arrives
    2. Window fires when watermark passes window end
    3. State kept for allowed_lateness period (for updates)
    4. State cleaned up after allowed_lateness expires
    """
    
    def __init__(
        self,
        allowed_lateness: timedelta,
        cleanup_interval: timedelta = timedelta(minutes=5)
    ):
        self.allowed_lateness = allowed_lateness
        self.cleanup_interval = cleanup_interval
        
        self.windows: Dict[str, Dict] = {}
        self.cleanup_timers: Dict[str, datetime] = {}
        self.last_cleanup = datetime.now(timezone.utc)
    
    def get_or_create_window(
        self,
        window_key: str,
        window_end: datetime
    ) -> Dict:
        """Get existing window state or create new."""
        
        if window_key not in self.windows:
            self.windows[window_key] = {
                "count": 0,
                "sum": 0,
                "created_at": datetime.now(timezone.utc),
                "window_end": window_end,
                "fired": False,
                "update_count": 0
            }
            
            # Set cleanup timer
            cleanup_time = window_end + self.allowed_lateness
            self.cleanup_timers[window_key] = cleanup_time
        
        return self.windows[window_key]
    
    def update_window(self, window_key: str, value: float):
        """Add value to window aggregate."""
        
        if window_key in self.windows:
            self.windows[window_key]["count"] += 1
            self.windows[window_key]["sum"] += value
            
            if self.windows[window_key]["fired"]:
                self.windows[window_key]["update_count"] += 1
    
    def fire_window(self, window_key: str) -> Optional[Dict]:
        """Fire window and return result."""
        
        if window_key not in self.windows:
            return None
        
        window = self.windows[window_key]
        window["fired"] = True
        
        return {
            "window_key": window_key,
            "count": window["count"],
            "sum": window["sum"],
            "avg": window["sum"] / window["count"] if window["count"] > 0 else 0,
            "is_final": False  # May still receive late updates
        }
    
    def maybe_cleanup(self, current_watermark: datetime):
        """
        Clean up expired window states.
        
        Called periodically to free memory.
        """
        
        now = datetime.now(timezone.utc)
        if now - self.last_cleanup < self.cleanup_interval:
            return
        
        self.last_cleanup = now
        
        windows_to_remove = []
        
        for window_key, cleanup_time in self.cleanup_timers.items():
            if current_watermark > cleanup_time:
                windows_to_remove.append(window_key)
        
        for window_key in windows_to_remove:
            # Emit final result before cleanup
            if window_key in self.windows:
                window = self.windows[window_key]
                logger.debug(
                    f"Cleaning up window {window_key}, "
                    f"received {window['update_count']} late updates"
                )
                
                del self.windows[window_key]
            
            del self.cleanup_timers[window_key]
        
        if windows_to_remove:
            logger.info(f"Cleaned up {len(windows_to_remove)} expired windows")
    
    def get_state_size(self) -> Dict:
        """Get current state size metrics."""
        return {
            "window_count": len(self.windows),
            "timer_count": len(self.cleanup_timers),
            "total_events": sum(w["count"] for w in self.windows.values())
        }

# late_data/reprocessing.py

"""
Reprocessing pipeline for late data and corrections.

Strategies:
1. Micro-batch reprocessing (hourly windows)
2. Daily batch reprocessing
3. On-demand backfill
"""

from dataclasses import dataclass
from datetime import datetime, date, timedelta, timezone
from typing import Dict, Any, List, Optional
import logging

logger = logging.getLogger(__name__)


@dataclass
class ReprocessingJob:
    """Defines a reprocessing job."""
    job_id: str
    start_time: datetime
    end_time: datetime
    reason: str
    status: str = "pending"
    created_at: datetime = None
    completed_at: datetime = None
    
    def __post_init__(self):
        if self.created_at is None:
            self.created_at = datetime.now(timezone.utc)


class ReprocessingManager:
    """
    Manages reprocessing of historical data.
    """
    
    def __init__(
        self,
        raw_data_source,
        aggregation_pipeline,
        output_sink
    ):
        self.source = raw_data_source
        self.pipeline = aggregation_pipeline
        self.sink = output_sink
        self.jobs: Dict[str, ReprocessingJob] = {}
    
    async def trigger_reprocessing(
        self,
        start_time: datetime,
        end_time: datetime,
        reason: str
    ) -> ReprocessingJob:
        """
        Trigger reprocessing for a time range.
        """
        
        job_id = f"reprocess_{start_time.strftime('%Y%m%d%H')}_{end_time.strftime('%Y%m%d%H')}"
        
        job = ReprocessingJob(
            job_id=job_id,
            start_time=start_time,
            end_time=end_time,
            reason=reason
        )
        
        self.jobs[job_id] = job
        
        logger.info(
            f"Triggered reprocessing job {job_id}: "
            f"{start_time} to {end_time}, reason: {reason}"
        )
        
        # Execute reprocessing
        await self._execute_reprocessing(job)
        
        return job
    
    async def _execute_reprocessing(self, job: ReprocessingJob):
        """Execute the reprocessing job."""
        
        try:
            job.status = "running"
            
            # Step 1: Read raw events for time range
            events = await self.source.read_events(
                start_time=job.start_time,
                end_time=job.end_time
            )
            
            logger.info(f"Read {len(events)} events for reprocessing")
            
            # Step 2: Run aggregation pipeline
            results = await self.pipeline.process_batch(events)
            
            # Step 3: Write results (replace existing)
            await self.sink.write_results(
                results,
                mode="overwrite",  # Replace, not append
                time_range=(job.start_time, job.end_time)
            )
            
            job.status = "completed"
            job.completed_at = datetime.now(timezone.utc)
            
            logger.info(
                f"Reprocessing job {job.job_id} completed, "
                f"processed {len(events)} events"
            )
            
        except Exception as e:
            job.status = "failed"
            logger.error(f"Reprocessing job {job.job_id} failed: {e}")
            raise
    
    async def schedule_daily_reprocessing(self, target_date: date):
        """
        Schedule end-of-day reprocessing to finalize numbers.
        
        Runs after allowed_lateness period to capture late data.
        """
        
        start = datetime.combine(target_date, datetime.min.time(), tzinfo=timezone.utc)
        end = datetime.combine(target_date + timedelta(days=1), datetime.min.time(), tzinfo=timezone.utc)
        
        return await self.trigger_reprocessing(
            start_time=start,
            end_time=end,
            reason="scheduled_daily_finalization"
        )


class IncrementalPatcher:
    """
    Applies incremental patches for late data.
    
    More efficient than full reprocessing for small corrections.
    """
    
    def __init__(self, result_store):
        self.store = result_store
    
    async def apply_late_events(
        self,
        late_events: List[Dict],
        aggregations: Dict[str, str]  # column -> aggregation type
    ) -> Dict[str, int]:
        """
        Apply late events as incremental patches.
        
        Only works for additive aggregations!
        """
        
        patches_by_window = {}
        
        for event in late_events:
            window_key = self._get_window_key(event)
            
            if window_key not in patches_by_window:
                patches_by_window[window_key] = {
                    "count": 0,
                    "updates": {}
                }
            
            patches_by_window[window_key]["count"] += 1
            
            for col, agg_type in aggregations.items():
                value = event.get(col, 0)
                
                if agg_type == "SUM":
                    patches_by_window[window_key]["updates"][col] = \
                        patches_by_window[window_key]["updates"].get(col, 0) + value
                
                elif agg_type == "COUNT":
                    patches_by_window[window_key]["updates"][col] = \
                        patches_by_window[window_key]["updates"].get(col, 0) + 1
                
                # Note: AVG, MEDIAN, DISTINCT require reprocessing!
        
        # Apply patches to store
        patched_count = 0
        for window_key, patch in patches_by_window.items():
            await self.store.apply_patch(window_key, patch["updates"])
            patched_count += patch["count"]
        
        return {
            "windows_patched": len(patches_by_window),
            "events_applied": patched_count
        }
    
    def _get_window_key(self, event: Dict) -> str:
        """Get window key from event."""
        event_time = datetime.fromisoformat(event["event_time"])
        return event_time.strftime("%Y-%m-%d-%H")

# late_data/versioning.py

"""
Versioned analytics results.

Instead of updating results in place, maintain versions:
- Each version has a timestamp
- Queries can request specific version or latest
- Audit trail of how numbers changed
"""

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


@dataclass
class ResultVersion:
    """A version of an analytics result."""
    version_id: int
    window_key: str
    result: Dict[str, Any]
    computed_at: datetime
    event_count: int
    is_final: bool
    reason: str  # "initial", "late_data", "correction", "reprocessing"


class VersionedResultStore:
    """
    Stores versioned analytics results.
    
    Benefits:
    1. Audit trail: See how numbers changed over time
    2. Debugging: Compare versions to find issues
    3. Consistency: Queries can pin to a version
    4. Rollback: Revert to previous version if needed
    """
    
    def __init__(self, storage_backend):
        self.storage = storage_backend
    
    async def store_result(
        self,
        window_key: str,
        result: Dict[str, Any],
        event_count: int,
        is_final: bool,
        reason: str
    ) -> ResultVersion:
        """
        Store a new version of a result.
        """
        
        # Get next version number
        existing = await self.storage.get_versions(window_key)
        next_version = len(existing) + 1
        
        version = ResultVersion(
            version_id=next_version,
            window_key=window_key,
            result=result,
            computed_at=datetime.now(timezone.utc),
            event_count=event_count,
            is_final=is_final,
            reason=reason
        )
        
        await self.storage.save_version(version)
        
        return version
    
    async def get_latest(self, window_key: str) -> Optional[ResultVersion]:
        """Get the most recent version."""
        versions = await self.storage.get_versions(window_key)
        return versions[-1] if versions else None
    
    async def get_final(self, window_key: str) -> Optional[ResultVersion]:
        """Get the final (official) version."""
        versions = await self.storage.get_versions(window_key)
        
        for v in reversed(versions):
            if v.is_final:
                return v
        
        return versions[-1] if versions else None
    
    async def get_at_time(
        self,
        window_key: str,
        as_of: datetime
    ) -> Optional[ResultVersion]:
        """Get the version that was current at a specific time."""
        
        versions = await self.storage.get_versions(window_key)
        
        # Find latest version computed before as_of time
        valid_versions = [
            v for v in versions if v.computed_at <= as_of
        ]
        
        return valid_versions[-1] if valid_versions else None
    
    async def get_version_history(
        self,
        window_key: str
    ) -> List[Dict]:
        """Get full version history for debugging."""
        
        versions = await self.storage.get_versions(window_key)
        
        history = []
        for i, v in enumerate(versions):
            entry = {
                "version": v.version_id,
                "computed_at": v.computed_at.isoformat(),
                "result": v.result,
                "event_count": v.event_count,
                "is_final": v.is_final,
                "reason": v.reason
            }
            
            # Add delta from previous version
            if i > 0:
                prev = versions[i-1]
                entry["delta"] = self._compute_delta(prev.result, v.result)
            
            history.append(entry)
        
        return history
    
    def _compute_delta(
        self,
        old_result: Dict,
        new_result: Dict
    ) -> Dict:
        """Compute difference between two result versions."""
        
        delta = {}
        
        for key in set(old_result.keys()) | set(new_result.keys()):
            old_val = old_result.get(key, 0)
            new_val = new_result.get(key, 0)
            
            if isinstance(old_val, (int, float)) and isinstance(new_val, (int, float)):
                diff = new_val - old_val
                if diff != 0:
                    delta[key] = {
                        "old": old_val,
                        "new": new_val,
                        "diff": diff,
                        "pct_change": (diff / old_val * 100) if old_val != 0 else None
                    }
        
        return delta


VERSION_HISTORY_EXAMPLE = """
VERSION HISTORY FOR WINDOW 2024-01-15-14

Version 1 (Initial):
  computed_at: 2024-01-15T15:00:00Z
  result: {revenue: 50000, orders: 500}
  event_count: 500
  is_final: false
  reason: initial

Version 2 (Late data):
  computed_at: 2024-01-15T16:30:00Z
  result: {revenue: 52500, orders: 525}
  event_count: 525
  is_final: false
  reason: late_data
  delta: {
    revenue: {old: 50000, new: 52500, diff: +2500, pct: +5%},
    orders: {old: 500, new: 525, diff: +25, pct: +5%}
  }

Version 3 (Final):
  computed_at: 2024-01-16T01:00:00Z
  result: {revenue: 53200, orders: 532}
  event_count: 532
  is_final: true
  reason: daily_finalization
  delta: {
    revenue: {old: 52500, new: 53200, diff: +700, pct: +1.3%},
    orders: {old: 525, new: 532, diff: +7, pct: +1.3%}
  }

This shows:
- Initial result was 5% low (mobile offline events)
- Daily finalization added 1.3% more (payment confirmations)
- Total correction: +6.4% from initial to final
"""

# late_data/production_pipeline.py

"""
Production late data handling pipeline.

Combines streaming for real-time with batch for accuracy.
"""

from dataclasses import dataclass, field
from datetime import datetime, timedelta, date, timezone
from typing import Dict, Any, List, Optional
import asyncio
import logging
from pyspark.sql import functions as F

logger = logging.getLogger(__name__)


@dataclass
class PipelineConfig:
    """Configuration for the late data pipeline."""
    
    # Watermark settings
    max_out_of_orderness: timedelta = timedelta(minutes=5)  # watermark lag
    allowed_lateness: timedelta = timedelta(minutes=5)      # updates after firing
    watermark_idle_timeout: timedelta = timedelta(minutes=1)
    
    # Window settings
    window_size: timedelta = timedelta(hours=1)
    
    # Reprocessing settings
    daily_reprocessing_hour: int = 2  # 2 AM
    reprocessing_lookback: timedelta = timedelta(days=1)
    
    # Output settings
    realtime_output: str = "druid://metrics"
    batch_output: str = "s3://warehouse/metrics"
    late_events_output: str = "kafka://late-events"


class LateDataPipeline:
    """
    Complete pipeline for handling late data in analytics.
    """
    
    def __init__(
        self,
        config: PipelineConfig,
        kafka_source,
        flink_env,
        spark_session,
        result_store: VersionedResultStore
    ):
        self.config = config
        self.kafka = kafka_source
        self.flink = flink_env
        self.spark = spark_session
        self.results = result_store
        
        # Components
        self.watermark_strategy = BoundedOutOfOrdernessWatermark(
            config.max_out_of_orderness
        )
        self.late_handler = LateDataHandler(
            LateDataConfig(
                allowed_lateness=config.allowed_lateness,
                action=LateDataAction.SIDE_OUTPUT,
                side_output_topic=config.late_events_output
            )
        )
        self.patcher = IncrementalPatcher(result_store)
        self.reprocessor = ReprocessingManager(
            raw_data_source=kafka_source,
            aggregation_pipeline=self,
            output_sink=result_store
        )
    
    async def start_streaming(self):
        """
        Start the streaming pipeline.
        
        This runs continuously, processing events as they arrive.
        """
        
        logger.info("Starting streaming pipeline")
        
        # Read from Kafka
        events = await self.kafka.read_stream("events")
        
        # Process events with watermarks
        async for event in events:
            await self._process_streaming_event(event)
    
    async def _process_streaming_event(self, event: Dict):
        """Process a single event in streaming mode."""
        
        event_time = datetime.fromisoformat(event["event_time"])
        
        # Update watermark
        watermark = self.watermark_strategy.on_event(event_time)
        
        # Process through late handler
        result = self.late_handler.process_event(
            event=event,
            watermark=watermark,
            aggregate_fn=self._aggregate_event
        )
        
        if result:
            # Store result version
            await self.results.store_result(
                window_key=result["window_key"],
                result=result["result"],
                event_count=result.get("event_count", 1),
                is_final=False,
                reason="late_data" if result["is_late_update"] else "streaming"
            )
    
    def _aggregate_event(
        self,
        current: Optional[Dict],
        event: Dict
    ) -> Dict:
        """Aggregate function for events."""
        
        if current is None:
            current = {
                "revenue": 0,
                "orders": 0,
                "items": 0
            }
        
        if event.get("event_type") == "order.placed":
            current["revenue"] += event.get("payload", {}).get("total", 0)
            current["orders"] += 1
            current["items"] += event.get("payload", {}).get("item_count", 0)
        
        return current
    
    async def run_daily_finalization(self, target_date: date):
        """
        Run daily finalization batch job.
        
        This reprocesses the full day to produce final numbers.
        """
        
        logger.info(f"Running daily finalization for {target_date}")
        
        start_time = datetime.combine(target_date, datetime.min.time(), tzinfo=timezone.utc)
        end_time = start_time + timedelta(days=1)
        
        # Read all events for the day from raw storage
        events_df = self.spark.read.parquet(
            f"s3://data-lake/events/date={target_date}"
        )
        
        # Aggregate by hour (F = pyspark.sql.functions)
        results = events_df.filter(
            "event_type = 'order.placed'"
        ).groupBy(
            F.date_trunc("hour", "event_time").alias("window_start")
        ).agg(
            F.sum("payload.total").alias("revenue"),
            F.count("*").alias("orders"),
            F.sum("payload.item_count").alias("items")
        ).collect()
        
        # Store final results
        for row in results:
            window_key = row["window_start"].strftime("%Y-%m-%d-%H")
            
            await self.results.store_result(
                window_key=window_key,
                result={
                    "revenue": row["revenue"],
                    "orders": row["orders"],
                    "items": row["items"]
                },
                event_count=row["orders"],
                is_final=True,
                reason="daily_finalization"
            )
        
        logger.info(
            f"Daily finalization complete for {target_date}, "
            f"finalized {len(results)} hourly windows"
        )
    
    async def apply_late_event_patches(self):
        """
        Process queued late events as patches.
        
        Runs periodically to update results with late data
        without full reprocessing.
        """
        
        late_events = self.late_handler.late_events
        
        if not late_events:
            return
        
        logger.info(f"Applying {len(late_events)} late event patches")
        
        # Turn raw late events into the measures _aggregate_event uses
        # (only orders contribute; raw events have no "revenue" field)
        contributions = [
            {
                "event_time": e["event_time"],
                "revenue": e.get("payload", {}).get("total", 0),
                "orders": 1,
                "items": e.get("payload", {}).get("item_count", 0),
            }
            for e in (le["event"] for le in late_events)
            if e.get("event_type") == "order.placed"
        ]
        
        # Apply as incremental patches (all three are additive)
        result = await self.patcher.apply_late_events(
            late_events=contributions,
            aggregations={
                "revenue": "SUM",
                "orders": "SUM",
                "items": "SUM"
            }
        )
        
        # Clear processed events
        self.late_handler.late_events = []
        
        logger.info(
            f"Applied patches to {result['windows_patched']} windows, "
            f"{result['events_applied']} events"
        )
    
    async def get_metrics_for_dashboard(
        self,
        window_key: str,
        version_type: str = "latest"  # "latest", "final", or ISO timestamp
    ) -> Dict:
        """
        Get metrics for dashboard display.
        
        Supports different version types:
        - latest: Most recent computation (may change)
        - final: Official finalized number (stable)
        - timestamp: As-of a specific time (for auditing)
        """
        
        if version_type == "latest":
            version = await self.results.get_latest(window_key)
        elif version_type == "final":
            version = await self.results.get_final(window_key)
        else:
            # Assume it's a timestamp
            as_of = datetime.fromisoformat(version_type)
            version = await self.results.get_at_time(window_key, as_of)
        
        if not version:
            return {"error": "No data found"}
        
        return {
            "window_key": window_key,
            "result": version.result,
            "computed_at": version.computed_at.isoformat(),
            "is_final": version.is_final,
            "version": version.version_id,
            "event_count": version.event_count,
            # Include staleness indicator
            "staleness": (datetime.now(timezone.utc) - version.computed_at).total_seconds()
        }

# late_data/dashboard_api.py

"""
API for dashboards to query late-data-aware metrics.
"""

from fastapi import FastAPI, Query
from datetime import datetime, date
from typing import Optional

app = FastAPI(title="Analytics API")


@app.get("/metrics/{metric_name}")
async def get_metric(
    metric_name: str,
    window: str = Query(..., description="Window key (e.g., 2024-01-15-14)"),
    version: str = Query(
        "latest",
        description="Version: 'latest', 'final', or ISO timestamp"
    ),
    include_history: bool = Query(
        False,
        description="Include version history"
    )
):
    """
    Get a metric value with late-data awareness.
    
    Returns:
    - result: The metric values
    - is_final: Whether this is the finalized number
    - staleness: Seconds since last computation
    - versions: (optional) Full version history
    """
    
    pipeline = get_pipeline()  # Dependency injection
    
    result = await pipeline.get_metrics_for_dashboard(
        window_key=window,
        version_type=version
    )
    
    if include_history:
        result["history"] = await pipeline.results.get_version_history(window)
    
    return result


@app.get("/metrics/{metric_name}/compare")
async def compare_versions(
    metric_name: str,
    window: str,
    version_a: str = "latest",
    version_b: str = "final"
):
    """
    Compare two versions of a metric.
    
    Useful for seeing real-time vs final discrepancy.
    """
    
    pipeline = get_pipeline()
    
    result_a = await pipeline.get_metrics_for_dashboard(window, version_a)
    result_b = await pipeline.get_metrics_for_dashboard(window, version_b)
    
    comparison = {
        "version_a": result_a,
        "version_b": result_b,
        "delta": {}
    }
    
    # Calculate delta
    for key in result_a.get("result", {}):
        val_a = result_a["result"].get(key, 0)
        val_b = result_b["result"].get(key, 0)
        
        if isinstance(val_a, (int, float)):
            comparison["delta"][key] = {
                "absolute": val_b - val_a,
                "percentage": ((val_b - val_a) / val_a * 100) if val_a else None
            }
    
    return comparison


@app.get("/data-quality/lateness")
async def get_lateness_metrics(
    start_date: date,
    end_date: date
):
    """
    Get data lateness metrics for monitoring.
    
    Shows how much data arrives late and how late.
    """
    
    pipeline = get_pipeline()
    handler = pipeline.late_handler
    
    metrics = handler.get_metrics()
    
    return {
        "period": {
            "start": start_date.isoformat(),
            "end": end_date.isoformat()
        },
        "on_time_rate": metrics["on_time_rate"],
        "late_rate": metrics["late_rate"],
        "late_accepted": metrics["late_accepted_count"],
        "late_dropped": metrics["late_dropped_count"],
        "recommendations": _get_recommendations(metrics)
    }


def _get_recommendations(metrics: dict) -> list:
    """Generate recommendations based on lateness metrics."""
    
    recommendations = []
    
    if metrics["late_rate"] > 0.10:  # >10% late
        recommendations.append({
            "severity": "high",
            "message": "High late data rate. Consider increasing allowed lateness.",
            "current_late_rate": f"{metrics['late_rate']*100:.1f}%"
        })
    
    if metrics["late_dropped_count"] > 0:
        recommendations.append({
            "severity": "medium",
            "message": f"{metrics['late_dropped_count']} events dropped as too late. "
                      "Review late event patterns."
        })
    
    return recommendations