"""
Data structures for waterfall enrichment.

These models track provider configuration, row-level enrichment state,
provenance events, and aggregated summary results.
"""

from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Dict, List, Optional, Set


class ProviderStatus(Enum):
    """Runtime status of a provider in the waterfall."""

    AVAILABLE = "available"
    UNAVAILABLE = "unavailable"  # Missing credentials
    RATE_LIMITED = "rate_limited"  # Hit rate limit
    CIRCUIT_OPEN = "circuit_open"  # Too many consecutive failures
    BUDGET_EXHAUSTED = "budget_exhausted"  # Hit budget/credit cap


@dataclass
class ProviderPlan:
    """Configuration for a single provider in the waterfall.

    Combines static configuration (from YAML) with runtime state tracking.
    """

    # Static configuration (from YAML)
    name: str
    priority: int
    fields: List[str]
    enabled: bool = True
    max_rows: Optional[int] = None
    max_calls: Optional[int] = None
    timeout_seconds: int = 60
    budget_credits: Optional[int] = None
    field_mapping: Dict[str, str] = field(default_factory=dict)

    # Runtime state (not persisted in config)
    status: ProviderStatus = ProviderStatus.AVAILABLE
    calls_made: int = 0
    rows_processed: int = 0
    credits_used: float = 0.0
    consecutive_failures: int = 0

    def can_provide_field(self, field_name: str) -> bool:
        """Check if this provider can fill the given field."""
        return field_name in self.fields or field_name in self.field_mapping.values()

    def get_output_column(self, provider_field: str) -> str:
        """Get the output column name for a provider field (applies mapping)."""
        return self.field_mapping.get(provider_field, provider_field)

    def is_available(self) -> bool:
        """Check if provider is available for use."""
        return self.enabled and self.status == ProviderStatus.AVAILABLE

    def check_limits(self) -> Optional[str]:
        """Check if any limits are exceeded. Returns reason string or None."""
        if self.max_rows is not None and self.rows_processed >= self.max_rows:
            return f"max_rows ({self.max_rows}) exceeded"
        if self.max_calls is not None and self.calls_made >= self.max_calls:
            return f"max_calls ({self.max_calls}) exceeded"
        if self.budget_credits is not None and self.credits_used >= self.budget_credits:
            return f"budget_credits ({self.budget_credits}) exceeded"
        return None

    def reset_runtime_state(self) -> None:
        """Reset runtime state for a new run."""
        self.status = ProviderStatus.AVAILABLE
        self.calls_made = 0
        self.rows_processed = 0
        self.credits_used = 0.0
        self.consecutive_failures = 0


@dataclass
class FieldRequest:
    """Tracks enrichment request for a specific field."""

    field_name: str
    output_column: str  # After mapping applied
    required: bool = True  # Part of stop condition
    providers_tried: List[str] = field(default_factory=list)
    filled_by: Optional[str] = None
    value: Optional[str] = None
    cache_hit: bool = False

    def is_filled(self) -> bool:
        """Check if field has a non-empty value."""
        if self.value is None:
            return False
        if isinstance(self.value, str) and not self.value.strip():
            return False
        return True


@dataclass
class RowState:
    """Tracks enrichment state for a single row."""

    row_index: int
    lookup_key: str  # Company name or domain used for lookup
    original_row: Dict[str, str]
    field_requests: Dict[str, FieldRequest] = field(default_factory=dict)

    # Computed state
    is_complete: bool = False  # Stop condition satisfied
    providers_used: Set[str] = field(default_factory=set)
    errors: List[str] = field(default_factory=list)

    def pending_fields(self) -> List[str]:
        """Return list of fields that still need to be filled."""
        return [f for f, req in self.field_requests.items() if not req.is_filled()]

    def filled_fields(self) -> List[str]:
        """Return list of fields that have been filled."""
        return [f for f, req in self.field_requests.items() if req.is_filled()]

    def get_field_value(self, field_name: str) -> Optional[str]:
        """Get the value of a field, or None if not filled."""
        req = self.field_requests.get(field_name)
        if req and req.is_filled():
            return req.value
        # Also check original row
        return self.original_row.get(field_name) or None

    def to_output_row(self) -> Dict[str, str]:
        """Merge enriched values into original row, adding metadata columns."""
        result = dict(self.original_row)

        # Apply enriched values
        for req in self.field_requests.values():
            if req.is_filled():
                result[req.output_column] = req.value

        # Add metadata columns
        enriched_by = sorted(self.providers_used)
        filled = [req.output_column for req in self.field_requests.values() if req.is_filled()]

        result["_enriched_by"] = ",".join(enriched_by) if enriched_by else ""
        result["_fields_filled"] = ",".join(sorted(filled)) if filled else ""
        result["_provider_errors"] = ";".join(self.errors) if self.errors else ""

        return result


@dataclass
class ProvenanceEvent:
    """Single enrichment attempt record for provenance sidecar."""

    timestamp: datetime
    row_index: int
    provider: str
    fields_requested: List[str]
    fields_filled: List[str]
    cache_hit: bool
    latency_ms: int
    success: bool
    error: Optional[str] = None
    credits_used: float = 0.0

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for JSONL output."""
        return {
            "ts": self.timestamp.isoformat(),
            "row": self.row_index,
            "provider": self.provider,
            "requested": self.fields_requested,
            "filled": self.fields_filled,
            "cache": self.cache_hit,
            "latency_ms": self.latency_ms,
            "ok": self.success,
            "error": self.error,
            "credits": self.credits_used,
        }


@dataclass
class ProviderResult:
    """Result from a single provider call."""

    success: bool
    data: Dict[str, Any] = field(default_factory=dict)
    error: Optional[str] = None
    retryable: bool = False
    latency_ms: int = 0
    credits_used: float = 0.0


@dataclass
class WaterfallSummary:
    """Aggregated results for AgenticResult output."""

    rows_total: int
    rows_enriched: int  # At least one field filled
    rows_complete: int  # Stop condition satisfied
    rows_failed: int  # All providers failed

    fields_requested: int
    fields_filled: int
    field_fill_rates: Dict[str, float]  # field -> % filled

    provider_stats: Dict[str, Dict[str, Any]]  # provider -> {calls, hits, misses, errors, credits}
    cache_stats: Dict[str, int]  # {hits, misses, expired}

    estimated_cost_usd: float
    total_latency_ms: int

    errors: List[str]  # Aggregated error messages
    warnings: List[str]

    def to_agentic_result(self, command: str, correlation_id: str, dry_run: bool = False) -> Dict[str, Any]:
        """Convert to AgenticResult-compatible dict."""
        return {
            "command": command,
            "status": "success" if self.rows_failed == 0 else "partial",
            "dry_run": dry_run,
            "result": {
                "rows_total": self.rows_total,
                "rows_enriched": self.rows_enriched,
                "rows_complete": self.rows_complete,
                "rows_failed": self.rows_failed,
                "fields_requested": self.fields_requested,
                "fields_filled": self.fields_filled,
                "field_fill_rates": self.field_fill_rates,
                "provider_stats": self.provider_stats,
                "cache_stats": self.cache_stats,
                "estimated_cost_usd": self.estimated_cost_usd,
                "total_latency_ms": self.total_latency_ms,
            },
            "correlation_id": correlation_id,
            "warnings": self.warnings,
            "errors": self.errors,
            "timestamp": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        }


@dataclass
class SafetyFlags:
    """Safety and control flags for waterfall execution."""

    dry_run: bool = False
    max_rows: Optional[int] = None
    sample: Optional[int] = None
    fail_fast: bool = False
    ignore_cache: bool = False
    budget_usd: Optional[float] = None  # Global spend limit


@dataclass
class CostEstimate:
    """Estimated cost breakdown for a waterfall run."""

    rows: int
    min_cost_usd: float
    max_cost_usd: float
    avg_cost_per_row: float
    provider_breakdown: Dict[str, Dict[str, float]]  # provider -> {cost_per_call, estimated_calls}
    cache_assumption: str  # e.g., "0% cache hit rate (worst case)"

    def __str__(self) -> str:
        if self.min_cost_usd == self.max_cost_usd:
            cost_str = f"${self.min_cost_usd:.2f}"
        else:
            cost_str = f"${self.min_cost_usd:.2f}-${self.max_cost_usd:.2f}"
        return f"Estimated cost: {cost_str} ({self.rows} rows × ~${self.avg_cost_per_row:.3f}/row)"


# Cost estimates per provider (USD per call) - used for estimating total cost
PROVIDER_COST_ESTIMATES: Dict[str, float] = {
    "foundrygraph": 0.0,  # Free (local)
    "cache": 0.0,  # Free (local)
    "apollo": 0.02,  # ~$0.01-0.03 per contact
    "pdl": 0.015,  # ~$0.01-0.02 per enrichment
    "clearbit": 0.02,  # ~$0.02 per enrich (adjust if you have contract rates)
    "apify": 0.075,  # ~$0.05-0.10 per run
    "zoominfo": 1.25,  # ~$0.50-2.00 per contact
}
