"""
Waterfall enrichment runner.

Orchestrates multi-provider enrichment with configurable ordering, stop conditions,
caching, rate limiting, circuit breaker, and provenance tracking.
"""

from __future__ import annotations

import csv
import logging
import random
import time
import unicodedata
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Dict, List, Set

from .circuit_breaker import CircuitBreaker
from .config import WaterfallConfig
from .models import (
    PROVIDER_COST_ESTIMATES,
    CostEstimate,
    FieldRequest,
    ProvenanceEvent,
    ProviderPlan,
    ProviderResult,
    ProviderStatus,
    RowState,
    SafetyFlags,
    WaterfallSummary,
)
from .provenance_writer import ProvenanceWriter
from .stop_conditions import StopCondition

logger = logging.getLogger(__name__)


class WaterfallError(Exception):
    """Error during waterfall enrichment."""

    pass


class WaterfallRunner:
    """Execute waterfall enrichment against a CSV file.

    The runner:
    1. Loads and validates the input CSV
    2. Applies safety flags (max_rows, sample)
    3. For each row, tries providers in priority order until stop condition is met
    4. Tracks provenance for each enrichment attempt
    5. Writes output CSV with metadata columns
    6. Writes provenance sidecar (JSONL)
    7. Returns summary statistics
    """

    def __init__(
        self,
        config: WaterfallConfig,
        safety_flags: SafetyFlags,
        company_column: str = "company",
    ):
        """Initialize waterfall runner.

        Args:
            config: Waterfall configuration
            safety_flags: Safety and control flags
            company_column: Column name for company lookup key
        """
        self.config = config
        self.safety = safety_flags
        self.company_column = company_column

        # Initialize components
        self.stop_condition = StopCondition(config.stop_condition_expression)
        self.circuit_breaker = CircuitBreaker(
            failure_threshold=5,
            reset_timeout_seconds=60.0,
        )
        self.provenance_events: List[ProvenanceEvent] = []

        # Rate limiting state
        self._rate_limit_tokens: Dict[str, float] = {}
        self._last_call_time: Dict[str, float] = {}

        # Provider dispatch map
        self._provider_dispatch: Dict[str, Callable] = {}
        self._init_provider_dispatch()

        # Stats tracking
        self._cache_hits = 0
        self._cache_misses = 0
        self._total_latency_ms = 0

    def _init_provider_dispatch(self) -> None:
        """Initialize provider dispatch functions."""
        # Import providers lazily to avoid circular imports
        try:
            from g_gremlin import apollo, zoominfo
            from g_gremlin import apify as apify_helpers
            from g_gremlin import pdl
            from g_gremlin import clearbit

            self._provider_dispatch["apollo"] = lambda company, fields: apollo.enrich_company(
                company, fields
            )
            self._provider_dispatch["zoominfo"] = lambda company, fields: zoominfo.enrich_company(
                company, fields
            )
            self._provider_dispatch["pdl"] = lambda company, fields: pdl.enrich_company(company, fields)
            self._provider_dispatch["clearbit"] = lambda company, fields: clearbit.enrich_company(
                company, fields
            )
            self._provider_dispatch["apify"] = lambda company, fields: self._apify_enrich(
                apify_helpers, company, fields
            )
        except ImportError as e:
            logger.warning(f"Some providers not available: {e}")

        try:
            from g_gremlin import foundrygraph

            self._provider_dispatch["foundrygraph"] = (
                lambda company, fields: self._foundrygraph_enrich(foundrygraph, company, fields)
            )
        except ImportError as e:
            logger.warning(f"FoundryGraph not available: {e}")

    def _apify_enrich(self, apify_module: Any, company: str, fields: List[str]) -> Dict[str, Any]:
        """Wrap Apify enrichment."""
        res = apify_module.enrich_company_sync(company, country="US")
        data = res.get("data", {}) if isinstance(res, dict) else {}
        if fields:
            return {k: data.get(k) for k in fields if k in data}
        return data

    def _foundrygraph_enrich(
        self, fg_module: Any, company: str, fields: List[str]
    ) -> Dict[str, Any]:
        """Wrap FoundryGraph enrichment."""
        # Try domain lookup first, then name search
        profile = None

        # If looks like a domain, try domain lookup
        if "." in company and " " not in company:
            profile = fg_module.lookup_by_domain(company)

        if not profile:
            # Try name search
            results = fg_module.search_by_name(company, limit=1)
            if results:
                profile = results[0]

        if not profile:
            return {}

        # Convert profile to dict
        if hasattr(profile, "__dict__"):
            data = vars(profile)
        else:
            data = dict(profile) if isinstance(profile, dict) else {}

        # Filter to requested fields
        if fields:
            return {k: data.get(k) for k in fields if k in data and data.get(k)}
        return data

    def estimate_cost(self, row_count: int) -> CostEstimate:
        """Estimate cost for enriching a given number of rows.

        Uses provider priority order and cost estimates to project costs.
        Assumes worst-case (0% cache hit rate) for max cost.

        Args:
            row_count: Number of rows to enrich

        Returns:
            CostEstimate with breakdown by provider
        """
        providers = self.config.get_providers_by_priority()
        if not providers:
            return CostEstimate(
                rows=row_count,
                min_cost_usd=0.0,
                max_cost_usd=0.0,
                avg_cost_per_row=0.0,
                provider_breakdown={},
                cache_assumption="No providers configured",
            )

        # Calculate costs based on provider order
        # Worst case: every row needs every provider (no stop condition satisfied early)
        # Best case: first provider satisfies stop condition for all rows
        provider_breakdown = {}
        total_max_cost = 0.0
        total_min_cost = 0.0

        for i, provider in enumerate(providers):
            cost_per_call = PROVIDER_COST_ESTIMATES.get(provider.name, 0.0)

            # Max: every row calls this provider (worst case)
            max_calls = row_count
            if provider.max_rows:
                max_calls = min(max_calls, provider.max_rows)
            if provider.max_calls:
                max_calls = min(max_calls, provider.max_calls)

            # Min: if this is the first paid provider, assume 50% need it
            # (optimistic: cache hits + earlier providers succeed)
            if cost_per_call > 0:
                min_calls = int(max_calls * 0.3) if i > 0 else int(max_calls * 0.5)
            else:
                min_calls = max_calls  # Free providers are always called

            provider_breakdown[provider.name] = {
                "cost_per_call": cost_per_call,
                "min_calls": min_calls,
                "max_calls": max_calls,
                "min_cost": min_calls * cost_per_call,
                "max_cost": max_calls * cost_per_call,
            }

            total_max_cost += max_calls * cost_per_call
            total_min_cost += min_calls * cost_per_call

        avg_cost = (total_min_cost + total_max_cost) / 2
        avg_per_row = avg_cost / row_count if row_count > 0 else 0.0

        return CostEstimate(
            rows=row_count,
            min_cost_usd=round(total_min_cost, 2),
            max_cost_usd=round(total_max_cost, 2),
            avg_cost_per_row=round(avg_per_row, 4),
            provider_breakdown=provider_breakdown,
            cache_assumption="0% cache hit rate (worst case for max, 50-70% for min)",
        )

    def _check_budget(self) -> bool:
        """Check if we've exceeded the global budget.

        Returns:
            True if budget is exceeded, False otherwise
        """
        if self.safety.budget_usd is None:
            return False

        total_spent = sum(p.credits_used for p in self.config.providers)
        return total_spent >= self.safety.budget_usd

    def run(
        self,
        input_csv: Path,
        output_csv: Path,
    ) -> WaterfallSummary:
        """Execute waterfall enrichment.

        Args:
            input_csv: Input CSV file path
            output_csv: Output CSV file path

        Returns:
            WaterfallSummary with statistics
        """
        run_id = str(uuid.uuid4())[:8]
        logger.info(f"Starting waterfall run {run_id} with config '{self.config.name}'")

        # Load and validate input
        rows = self._load_csv(input_csv)
        if not rows:
            return self._build_empty_summary()

        # Apply safety limits
        rows = self._apply_safety_limits(rows)

        # Check provider availability
        self._check_provider_availability()

        # Initialize row states
        row_states = [self._init_row_state(i, row) for i, row in enumerate(rows)]

        # Initialize provenance writer
        provenance_enabled = self.config.output.provenance_mode == "sidecar"
        with ProvenanceWriter(
            output_csv, enabled=provenance_enabled and not self.safety.dry_run
        ) as provenance:
            provenance.start_run(run_id, self.config.name)

            # Process rows in chunks
            chunk_size = 500
            for i in range(0, len(row_states), chunk_size):
                chunk = row_states[i : i + chunk_size]
                self._process_chunk(chunk)

                # Flush provenance for this chunk
                for event in self.provenance_events:
                    provenance.add_event(event)
                self.provenance_events.clear()

            # Build summary before writing
            summary = self._build_summary(row_states)

            # Write output
            if not self.safety.dry_run:
                self._write_output(row_states, output_csv)
                provenance.finish_run(summary)

        logger.info(
            f"Waterfall run {run_id} complete: {summary.rows_enriched}/{summary.rows_total} enriched"
        )
        return summary

    def _load_csv(self, csv_path: Path) -> List[Dict[str, str]]:
        """Load input CSV file."""
        if not csv_path.exists():
            raise WaterfallError(f"Input CSV not found: {csv_path}")

        with csv_path.open(newline="", encoding="utf-8-sig") as f:
            reader = list(csv.DictReader(f))

        if not reader:
            logger.warning(f"Empty CSV file: {csv_path}")
            return []

        if self.company_column not in reader[0]:
            raise WaterfallError(
                f"Company column '{self.company_column}' not found in CSV. "
                f"Available columns: {list(reader[0].keys())}"
            )

        return reader

    def _apply_safety_limits(self, rows: List[Dict[str, str]]) -> List[Dict[str, str]]:
        """Apply max_rows and sample limits."""
        if self.safety.sample:
            sample_size = min(self.safety.sample, len(rows))
            rows = random.sample(rows, sample_size)
            logger.info(f"Sampled {sample_size} rows from {len(rows)}")

        if self.safety.max_rows and len(rows) > self.safety.max_rows:
            logger.info(f"Limiting to {self.safety.max_rows} rows (from {len(rows)})")
            rows = rows[: self.safety.max_rows]

        return rows

    def _check_provider_availability(self) -> None:
        """Check which providers are available and mark unavailable ones."""
        for provider in self.config.providers:
            if provider.name not in self._provider_dispatch:
                provider.status = ProviderStatus.UNAVAILABLE
                logger.warning(
                    f"Provider '{provider.name}' not available (no dispatch function)"
                )

    def _init_row_state(self, index: int, row: Dict[str, str]) -> RowState:
        """Initialize row state for processing."""
        company = row.get(self.company_column, "").strip()
        lookup_key = self._normalize_company(company)

        # Determine which fields need to be filled
        # Combine all provider fields and stop condition fields
        all_fields: Set[str] = set()
        for provider in self.config.providers:
            all_fields.update(provider.fields)
        all_fields.update(self.stop_condition.required_fields)

        # Initialize field requests for missing fields
        field_requests = {}
        for field in all_fields:
            # Check if field already has a value in the row
            existing = row.get(field, "").strip()
            if existing:
                # Field already filled in source data
                field_requests[field] = FieldRequest(
                    field_name=field,
                    output_column=field,
                    value=existing,
                    filled_by="source",
                )
            else:
                field_requests[field] = FieldRequest(
                    field_name=field,
                    output_column=field,
                )

        state = RowState(
            row_index=index,
            lookup_key=lookup_key,
            original_row=dict(row),
            field_requests=field_requests,
        )

        # Check if already complete from source data
        state.is_complete = self.stop_condition.evaluate_row_state(state)

        return state

    def _normalize_company(self, name: str) -> str:
        """Normalize company name for cache key and lookup."""
        # NFKC normalization for Unicode
        normalized = unicodedata.normalize("NFKC", name)
        return normalized.strip().lower()

    def _process_chunk(self, rows: List[RowState]) -> None:
        """Process a chunk of rows."""
        for row in rows:
            if row.is_complete:
                continue

            if not row.lookup_key:
                row.errors.append("empty_lookup_key")
                continue

            # Check global budget before processing each row
            if self._check_budget():
                logger.warning(
                    f"Budget limit ${self.safety.budget_usd:.2f} reached, "
                    f"stopping at row {row.row_index}"
                )
                row.errors.append("budget_exceeded")
                break

            self._process_row(row)

    def _process_row(self, row: RowState) -> None:
        """Process a single row through the waterfall."""
        pending = row.pending_fields()
        if not pending:
            row.is_complete = self.stop_condition.evaluate_row_state(row)
            return

        # Get providers sorted by priority
        providers = self.config.get_providers_by_priority()

        for provider in providers:
            if row.is_complete:
                break

            if not provider.is_available():
                continue

            # Check if provider can fill any pending fields
            provider_fields = set(provider.fields)
            needed_from_provider = [f for f in pending if f in provider_fields]
            if not needed_from_provider:
                continue

            # Check provider limits
            limit_reason = provider.check_limits()
            if limit_reason:
                if provider.status != ProviderStatus.BUDGET_EXHAUSTED:
                    logger.warning(f"Provider {provider.name}: {limit_reason}")
                    if "budget" in limit_reason.lower():
                        provider.status = ProviderStatus.BUDGET_EXHAUSTED
                continue

            # Check circuit breaker
            if self.circuit_breaker.is_open(provider.name):
                continue

            # Try cache first (unless ignore_cache is set)
            if not self.safety.ignore_cache:
                cache_hits = self._check_cache(provider, row, needed_from_provider)
                if cache_hits:
                    self._apply_values(row, provider.name, cache_hits, cache_hit=True)
                    pending = row.pending_fields()
                    row.is_complete = self.stop_condition.evaluate_row_state(row)
                    if row.is_complete:
                        break

            # Remaining fields after cache
            pending = row.pending_fields()
            needed_from_provider = [f for f in pending if f in provider_fields]
            if not needed_from_provider:
                continue

            # Call provider
            result = self._call_provider(provider, row, needed_from_provider)

            # Record provenance
            event = ProvenanceEvent(
                timestamp=datetime.now(timezone.utc),
                row_index=row.row_index,
                provider=provider.name,
                fields_requested=needed_from_provider,
                fields_filled=list(result.data.keys()) if result.success else [],
                cache_hit=False,
                latency_ms=result.latency_ms,
                success=result.success,
                error=result.error,
                credits_used=result.credits_used,
            )
            self.provenance_events.append(event)

            if result.success:
                self._apply_values(row, provider.name, result.data, cache_hit=False)
                self._cache_values(provider, row, result.data)
                provider.rows_processed += 1
                self.circuit_breaker.record_success(provider.name)
            else:
                row.errors.append(f"{provider.name}:{result.error}")
                self.circuit_breaker.record_failure(provider.name)

                if self.safety.fail_fast:
                    raise WaterfallError(f"Provider {provider.name} failed: {result.error}")

            # Update completion status
            pending = row.pending_fields()
            row.is_complete = self.stop_condition.evaluate_row_state(row)

    def _check_cache(
        self, provider: ProviderPlan, row: RowState, fields: List[str]
    ) -> Dict[str, str]:
        """Check cache for field values."""
        from g_gremlin.enrichment_cache import get_cache

        hits = {}
        for field in fields:
            entry = get_cache(provider.name, row.lookup_key, field)
            if entry:
                hits[field] = entry.value
                self._cache_hits += 1
            else:
                self._cache_misses += 1

        return hits

    def _cache_values(
        self, provider: ProviderPlan, row: RowState, values: Dict[str, Any]
    ) -> None:
        """Store values in cache."""
        from g_gremlin.enrichment_cache import set_cache

        ttl = self.config.cache.ttl_seconds
        company_raw = row.original_row.get(self.company_column, "")

        for field, value in values.items():
            if value:
                set_cache(
                    provider=provider.name,
                    company_raw=company_raw,
                    company_norm=row.lookup_key,
                    field=field,
                    value=str(value),
                    confidence=1.0,
                    ttl_seconds=ttl,
                )

    def _apply_values(
        self,
        row: RowState,
        provider_name: str,
        values: Dict[str, Any],
        cache_hit: bool,
    ) -> None:
        """Apply enriched values to row state."""
        for field, value in values.items():
            if field not in row.field_requests:
                continue

            req = row.field_requests[field]
            str_value = str(value) if value else ""

            # First-fill-wins, except empty->non-empty overwrites
            if not req.is_filled() and str_value.strip():
                req.value = str_value
                req.filled_by = provider_name
                req.cache_hit = cache_hit
                row.providers_used.add(provider_name)

    def _call_provider(
        self, provider: ProviderPlan, row: RowState, fields: List[str]
    ) -> ProviderResult:
        """Call a provider with rate limiting and error handling."""
        # Rate limiting
        self._apply_rate_limit(provider.name)

        dispatch_fn = self._provider_dispatch.get(provider.name)
        if not dispatch_fn:
            return ProviderResult(
                success=False,
                error=f"No dispatch function for {provider.name}",
            )

        try:
            start = time.monotonic()
            company = row.original_row.get(self.company_column, "")
            data = dispatch_fn(company, fields)
            latency_ms = int((time.monotonic() - start) * 1000)
            self._total_latency_ms += latency_ms

            provider.calls_made += 1

            # Estimate credits used
            credits = PROVIDER_COST_ESTIMATES.get(provider.name, 0.0)
            provider.credits_used += credits

            # Filter to requested fields and non-empty values
            filtered = {k: v for k, v in data.items() if k in fields and v}

            return ProviderResult(
                success=True,
                data=filtered,
                latency_ms=latency_ms,
                credits_used=credits,
            )

        except Exception as e:
            provider.consecutive_failures += 1
            return ProviderResult(
                success=False,
                error=str(e),
                retryable=True,
            )

    def _apply_rate_limit(self, provider_name: str) -> None:
        """Apply rate limiting for a provider."""
        rate_per_minute = self.config.rate_limit.requests_per_minute
        min_interval = 60.0 / rate_per_minute

        last_call = self._last_call_time.get(provider_name, 0.0)
        now = time.monotonic()
        elapsed = now - last_call

        if elapsed < min_interval:
            sleep_time = min_interval - elapsed
            time.sleep(sleep_time)

        self._last_call_time[provider_name] = time.monotonic()

    def _write_output(self, row_states: List[RowState], output_path: Path) -> None:
        """Write enriched rows to output CSV."""
        if not row_states:
            return

        output_path.parent.mkdir(parents=True, exist_ok=True)

        # Build output rows
        output_rows = [row.to_output_row() for row in row_states]

        # Determine fieldnames (original + new fields + metadata)
        fieldnames = list(row_states[0].original_row.keys())
        for row in output_rows:
            for key in row:
                if key not in fieldnames:
                    fieldnames.append(key)

        with output_path.open("w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            writer.writerows(output_rows)

        logger.info(f"Wrote {len(output_rows)} rows to {output_path}")

    def _build_summary(self, row_states: List[RowState]) -> WaterfallSummary:
        """Build summary statistics from row states."""
        rows_total = len(row_states)
        rows_enriched = sum(1 for r in row_states if r.providers_used)
        rows_complete = sum(1 for r in row_states if r.is_complete)
        rows_failed = sum(1 for r in row_states if r.errors and not r.providers_used)

        # Field fill rates
        all_fields: Set[str] = set()
        for r in row_states:
            all_fields.update(r.field_requests.keys())

        field_fill_rates = {}
        fields_filled = 0
        for field in all_fields:
            filled_count = sum(
                1 for r in row_states if field in r.field_requests and r.field_requests[field].is_filled()
            )
            field_fill_rates[field] = filled_count / rows_total if rows_total > 0 else 0.0
            fields_filled += filled_count

        # Provider stats
        provider_stats = {}
        for provider in self.config.providers:
            provider_stats[provider.name] = {
                "calls": provider.calls_made,
                "rows_processed": provider.rows_processed,
                "credits_used": provider.credits_used,
                "status": provider.status.value,
                "consecutive_failures": provider.consecutive_failures,
            }

        # Cache stats
        cache_stats = {
            "hits": self._cache_hits,
            "misses": self._cache_misses,
        }

        # Estimated cost
        estimated_cost = sum(p.credits_used for p in self.config.providers)

        # Collect errors
        all_errors = []
        for r in row_states:
            all_errors.extend(r.errors)

        return WaterfallSummary(
            rows_total=rows_total,
            rows_enriched=rows_enriched,
            rows_complete=rows_complete,
            rows_failed=rows_failed,
            fields_requested=len(all_fields) * rows_total,
            fields_filled=fields_filled,
            field_fill_rates=field_fill_rates,
            provider_stats=provider_stats,
            cache_stats=cache_stats,
            estimated_cost_usd=estimated_cost,
            total_latency_ms=self._total_latency_ms,
            errors=list(set(all_errors))[:10],  # Limit to 10 unique errors
            warnings=[],
        )

    def _build_empty_summary(self) -> WaterfallSummary:
        """Build summary for empty input."""
        return WaterfallSummary(
            rows_total=0,
            rows_enriched=0,
            rows_complete=0,
            rows_failed=0,
            fields_requested=0,
            fields_filled=0,
            field_fill_rates={},
            provider_stats={},
            cache_stats={"hits": 0, "misses": 0},
            estimated_cost_usd=0.0,
            total_latency_ms=0,
            errors=[],
            warnings=["Empty input CSV"],
        )
