"""
Enhanced RunStats with proper progress emission.
"""

import time
import logging
import multiprocessing as mp
from dataclasses import dataclass, field
from typing import Optional
from .progress_event import ProgressEvent, ProgressPhase

log = logging.getLogger(__name__)


@dataclass
class RunStats:
    """
    Tracks processing statistics and emits progress events.

    Features:
    - Time-based emission (not just count-based)
    - Tracks both items and comparisons
    - Heartbeat support for long operations
    - Final emission guarantee
    - Dual counter display (items + blocks)
    - Progress basis tagging
    """

    # Configuration
    total_items_to_process: int
    job_id: Optional[str] = None  # ADDED: job_id for final event identification
    emit_interval_s: float = 0.1
    emit_threshold_pct: float = 0.01  # 1% progress threshold
    heartbeat_interval_s: float = 2.0
    console_emit_interval_s: float = 2.0  # Rate-limit console logging

    # State tracking
    total_comparisons: int = 0
    items_processed_count: int = 0
    blocks_processed_count: int = 0  # NEW: Track blocks separately
    total_blocks: int = 0  # NEW: Total blocks to process
    start_offset: int = 0
    t_overall_start: float = field(default_factory=time.perf_counter)

    # Progress mode tracking
    progress_basis: str = "items"  # NEW: "items" or "blocks"

    # Emission tracking
    _last_emit_time: float = field(default_factory=time.perf_counter)
    _last_heartbeat_time: float = field(default_factory=time.perf_counter)
    _last_console_emit_time: float = field(
        default_factory=time.perf_counter
    )  # NEW: Console rate limiting
    _last_emitted_items: int = 0
    _last_emitted_blocks: int = 0  # NEW: Track emitted blocks
    _last_emitted_comparisons: int = 0
    _last_progress_q: Optional[mp.Queue] = None
    _current_phase: str = ProgressPhase.INIT.value

    def __post_init__(self):
        """Initialize and emit initial progress event."""
        # Add this line:
        import threading

        self._lock = threading.Lock()

        if self._last_progress_q:
            self._emit_event(
                ProgressEvent(
                    phase=ProgressPhase.INIT.value,
                    total_items=self.total_items_to_process,
                    message="Initializing processing...",
                    timestamp_ns=time.perf_counter_ns(),
                    extras={},
                )
            )

    def update(
        self,
        comparisons_in_chunk: int,
        items_processed_in_chunk: int,
        progress_q: Optional[mp.Queue] = None,
    ) -> None:
        """
        Update statistics and emit progress if needed.

        Emission triggers:
        1. Time elapsed > emit_interval_s
        2. Progress > emit_threshold_pct (for either items OR comparisons)
        3. First update
        4. Reached 100% completion
        """
        # Update totals
        self.total_comparisons += comparisons_in_chunk
        self.items_processed_count += items_processed_in_chunk

        # Store queue reference
        if progress_q:
            self._last_progress_q = progress_q

        # Check if we should emit
        if self._should_emit():
            self._emit_progress()

    def _should_emit(self) -> bool:
        """Determine if progress should be emitted."""
        current_time = time.perf_counter()

        # Time-based emission
        if current_time - self._last_emit_time >= self.emit_interval_s:
            return True

        # First emission
        if self._last_emitted_items == 0 and self._last_emitted_comparisons == 0:
            return True

        # Percentage-based emission (items)
        if self.total_items_to_process > 0:
            old_pct = self._last_emitted_items / self.total_items_to_process
            new_pct = self.items_processed_count / self.total_items_to_process
            if new_pct - old_pct >= self.emit_threshold_pct:
                return True

        # Percentage-based emission (comparisons)
        # Estimate total comparisons based on items if possible
        if self.total_comparisons > self._last_emitted_comparisons:
            # For significant comparison increases, always emit
            comp_delta = self.total_comparisons - self._last_emitted_comparisons
            if comp_delta > 10000:  # Significant work done
                return True

        # Completion check
        if (
            self.total_items_to_process > 0
            and self.items_processed_count >= self.total_items_to_process
        ):
            return True

        return False

    def emit_heartbeat(self, message: str = "Processing...") -> None:
        """Emit a heartbeat event if enough time has passed."""
        current_time = time.perf_counter()

        if (
            self._last_progress_q
            and current_time - self._last_heartbeat_time >= self.heartbeat_interval_s
        ):
            self._emit_event(
                ProgressEvent(
                    phase=ProgressPhase.HEARTBEAT.value,
                    total_items=self.total_items_to_process,
                    message=message,
                    timestamp_ns=time.perf_counter_ns(),
                    extras={"elapsed_s": current_time - self.t_overall_start},
                )
            )
            self._last_heartbeat_time = current_time

    def set_phase(self, phase: str) -> None:
        """Update the current processing phase with protection against invalid transitions."""
        current_phase = self._current_phase

        # Prevent going back from FINAL to RUN
        if (
            current_phase == ProgressPhase.FINAL.value
            and phase == ProgressPhase.RUN.value
        ):
            log.warning(f"Blocked invalid phase transition: {current_phase} -> {phase}")
            return

        # Prevent premature FINAL transition
        if (
            phase == ProgressPhase.FINAL.value
            and self.items_processed_count < self.total_items_to_process * 0.95
        ):
            log.warning(
                f"Blocked premature FINAL transition: only {self.items_processed_count}/{self.total_items_to_process} processed"
            )
            return

        self._current_phase = phase
        log.debug(f"Phase transition: {current_phase} -> {phase}")

        if phase == ProgressPhase.RUN.value:
            # Emit a phase transition event
            self._emit_progress(force=True)

    def set_progress_basis(self, basis: str, total_blocks: int = None) -> None:
        """
        Switch progress tracking basis between 'items' and 'blocks'.
        Called when switching to parallel processing.
        """
        old_basis = self.progress_basis
        self.progress_basis = basis

        if total_blocks is not None:
            self.total_blocks = total_blocks

        # Log the mode switch with carryover information
        log.info(f"Progress tracking switched from basis={old_basis} to basis={basis}")

        # Emit carryover breadcrumb for transition clarity
        if old_basis != basis:
            log.info(
                f"[progress-transition] Carryover totals: "
                f"items={self.items_processed_count:,}/{self.total_items_to_process:,} "
                f"blocks={self.blocks_processed_count:,}/{self.total_blocks:,} "
                f"comparisons={self.total_comparisons:,}"
            )

        # Force immediate progress emission with new basis
        self._emit_progress(force=True, log_to_console=True)

    def _emit_progress(self, force: bool = False, log_to_console: bool = False) -> None:
        """Emit accumulated progress with dual counter display."""
        if not self._last_progress_q:
            return

        current_time = time.perf_counter()

        # Calculate deltas
        items_delta = self.items_processed_count - self._last_emitted_items
        blocks_delta = self.blocks_processed_count - self._last_emitted_blocks
        comparisons_delta = self.total_comparisons - self._last_emitted_comparisons

        # Skip if no progress (unless forced)
        if (
            not force
            and items_delta == 0
            and comparisons_delta == 0
            and blocks_delta == 0
        ):
            return

        # Determine appropriate phase - with protection against premature finalization
        phase = self._current_phase
        if (
            self.total_items_to_process > 0
            and self.items_processed_count >= self.total_items_to_process
            and self.items_processed_count >= self.total_items_to_process * 0.95
        ):  # ADDED: 95% threshold check
            phase = ProgressPhase.FINAL.value

        # Calculate progress percentage based on current basis
        if self.progress_basis == "blocks" and self.total_blocks > 0:
            progress_pct = self.blocks_processed_count / self.total_blocks * 100
        elif self.total_items_to_process > 0:
            progress_pct = (
                self.items_processed_count / self.total_items_to_process * 100
            )
        else:
            progress_pct = 0

        # Create enhanced progress message with dual counters
        if phase == ProgressPhase.FINAL.value:
            message = "Processing complete"
        else:
            # Format: progress[basis=blocks] 63% (items 78,xxx/199,061; blocks 2,071/3,280)
            message = (
                f"progress[basis={self.progress_basis}] {progress_pct:.1f}% "
                f"(items {self.items_processed_count:,}/{self.total_items_to_process:,}; "
                f"blocks {self.blocks_processed_count:,}/{self.total_blocks:,})"
            )

        # Console logging with rate limiting
        should_log_console = (
            log_to_console
            or phase == ProgressPhase.FINAL.value
            or (
                current_time - self._last_console_emit_time
                >= self.console_emit_interval_s
            )
        )

        if should_log_console:
            log.info(message)
            self._last_console_emit_time = current_time

        # Create base extras dictionary
        extras = {
            "elapsed_s": current_time - self.t_overall_start,
            "items_rate": (
                self.items_processed_count / (current_time - self.t_overall_start)
                if current_time > self.t_overall_start
                else 0
            ),
            "progress_basis": self.progress_basis,  # Include basis in extras
            "blocks_processed": self.blocks_processed_count,
            "total_blocks": self.total_blocks,
        }

        # Include job_id in extras for FINAL phase events
        if phase == ProgressPhase.FINAL.value and self.job_id:
            extras["op_id"] = self.job_id

        # Emit event
        event = ProgressEvent(
            items_delta=items_delta,
            comparisons_delta=comparisons_delta,
            total_items=self.total_items_to_process,
            phase=phase,
            message=message,
            timestamp_ns=time.perf_counter_ns(),
            extras=extras,
        )

        self._emit_event(event)

        # Update tracking
        self._last_emitted_items = self.items_processed_count
        self._last_emitted_blocks = self.blocks_processed_count
        self._last_emitted_comparisons = self.total_comparisons
        self._last_emit_time = current_time

    def _emit_event(self, event: ProgressEvent) -> None:
        """Send event to progress queue."""
        try:
            # We now ONLY send the modern ProgressEvent object.
            self._last_progress_q.put(event)
        except Exception as e:
            log.error(f"Failed to emit progress event: {e}")

    def update_blocks(
        self,
        blocks_processed: int,
        total_blocks: int,
        items_in_chunk: int,  # ADDED: new parameter for actual items processed
        comparisons_in_chunk: int = 0,
        progress_q: Optional[mp.Queue] = None,
    ) -> None:
        """
        Update statistics based on blocks processed.
        Uses items_in_chunk for accurate progress tracking.
        """
        # Update total blocks if provided
        if total_blocks > 0:
            self.total_blocks = total_blocks

        # Update block counter
        self.blocks_processed_count = blocks_processed

        # Call regular update with the actual number of items (rows) processed
        self.update(comparisons_in_chunk, items_in_chunk, progress_q)

    def force_final_emit(self) -> None:
        """Force emission of any remaining progress and ensure it's marked as FINAL."""
        if not self._last_progress_q:
            return

        log.info(f"force_final_emit called with job_id: {self.job_id}")

        # CRITICAL FIX: Don't force FINAL phase if we're not done processing
        if (
            self.items_processed_count < self.total_items_to_process * 0.95
        ):  # 95% threshold
            log.warning(
                f"force_final_emit called prematurely: only {self.items_processed_count}/{self.total_items_to_process} items processed"
            )
            # Emit current progress but don't change phase
            self._emit_progress(force=True)
            return

        # Check if a final event for this job has already been emitted with all progress
        if (
            self._current_phase == ProgressPhase.FINAL.value
            and self.items_processed_count == self._last_emitted_items
            and self.total_comparisons == self._last_emitted_comparisons
        ):
            log.debug(
                "Final event with all progress already sent, skipping redundant emit."
            )
            return

        # Only set FINAL if we're truly done
        self._current_phase = ProgressPhase.FINAL.value
        extras = {}
        if self.job_id:
            extras["op_id"] = self.job_id

        # Check for any unemitted progress
        items_remaining = self.items_processed_count - self._last_emitted_items
        comps_remaining = self.total_comparisons - self._last_emitted_comparisons

        if items_remaining > 0 or comps_remaining > 0:
            log.info(
                f"Emitting FINAL progress with deltas - job_id in extras: {extras.get('op_id')}"
            )
            self._emit_progress(force=True)
        else:
            # Even if no delta, emit a final event to signal completion
            log.info(
                f"Emitting FINAL event with no deltas - job_id in extras: {extras.get('op_id')}"
            )
            self._emit_event(
                ProgressEvent(
                    phase=ProgressPhase.FINAL.value,
                    total_items=self.total_items_to_process,
                    message="Processing complete",
                    timestamp_ns=time.perf_counter_ns(),
                    extras=extras,
                )
            )
