# SPDX-License-Identifier: Apache-2.0
# Standard
from typing import Optional, Tuple
import time

# First Party
from lmcache.logging import init_logger
from lmcache.v1.cache_controller.utils import (
    FullSyncState,
    RegistryTree,
    WorkerNode,
    WorkerSyncInfo,
)

logger = init_logger(__name__)


class FullSyncTracker:
    """
    Tracks full sync state for all workers.

    This class manages the state of full sync operations, including:
    - Tracking which workers need full sync
    - Monitoring sync progress
    - Handling sync timeout
    - Determining when freeze mode can be exited
    """

    def __init__(
        self,
        registry_tree: RegistryTree,
        completion_threshold: float = 0.8,
        sync_timeout_s: float = 300.0,
    ):
        """
        Initialize the FullSyncTracker.

        Args:
            registry_tree: The registry tree containing worker nodes
            completion_threshold: Percentage of workers that need to complete
                                  sync before others can exit freeze mode (default: 80%)
            sync_timeout_s: Timeout in seconds for a single worker's sync
                                  (default: 300s)
        """
        self.registry_tree = registry_tree
        self.completion_threshold = completion_threshold
        self.sync_timeout_s = sync_timeout_s

        # Flag to indicate if controller just restarted and needs full sync
        self._need_full_sync_all = True

    def _get_sync_info(
        self, instance_id: str, worker_id: int
    ) -> Optional[WorkerSyncInfo]:
        """Get sync info for a worker from the registry tree."""
        worker_node = self.registry_tree.get_worker(instance_id, worker_id)
        if worker_node is None:
            return None
        return worker_node.sync_info

    def _set_sync_info(
        self, instance_id: str, worker_id: int, sync_info: Optional[WorkerSyncInfo]
    ) -> bool:
        """Set sync info for a worker. Returns True if successful."""
        worker_node = self.registry_tree.get_worker(instance_id, worker_id)
        if worker_node is None:
            return False
        worker_node.sync_info = sync_info
        return True

    def set_need_full_sync_all(self, need: bool) -> None:
        """Set whether all workers need full sync (e.g., after controller restart)"""
        self._need_full_sync_all = need
        logger.info("Set need_full_sync_all to %s", need)

    def _get_all_workers_cached(
        self, timeout_seconds: Optional[float] = None
    ) -> list[tuple[str, WorkerNode]]:
        """Get all registered workers from the registry tree."""
        return self.registry_tree.get_all_worker_nodes_cached(timeout_seconds)

    def should_request_full_sync(
        self, instance_id: str, worker_id: int
    ) -> Tuple[bool, Optional[str]]:
        """
        Check if a worker should perform full sync.

        Returns:
            Tuple of (need_sync, reason)
        """
        # Case 1: Controller just restarted, all workers need sync
        if self._need_full_sync_all:
            sync_info = self._get_sync_info(instance_id, worker_id)
            if sync_info is None or sync_info.state not in (
                FullSyncState.SYNCING,
                FullSyncState.COMPLETED,
            ):
                return True, "controller_restart"

        # Case 2: Worker sync failed/timeout, needs retry
        sync_info = self._get_sync_info(instance_id, worker_id)
        if sync_info is not None and sync_info.state == FullSyncState.FAILED:
            return True, "sync_failed_retry"

        return False, None

    def is_worker_syncing(self, instance_id: str, worker_id: int) -> bool:
        """
        Check if a worker is currently in sync state.

        When a worker is syncing, incremental events should be discarded.
        """
        sync_info = self._get_sync_info(instance_id, worker_id)
        if sync_info is None:
            return False
        return sync_info.state == FullSyncState.SYNCING

    def get_sync_id(self, instance_id: str, worker_id: int) -> Optional[str]:
        """
        Get sync ID for a worker.

        Returns sync ID if worker is syncing, None otherwise.
        """
        sync_info = self._get_sync_info(instance_id, worker_id)
        if sync_info is None:
            return None
        return sync_info.sync_id

    def start_sync(
        self,
        instance_id: str,
        worker_id: int,
        sync_id: str,
        total_keys: int,
        batch_count: int,
    ) -> bool:
        """
        Start sync for a worker.

        Returns:
            True if sync started successfully, False otherwise
        """
        report_id = (instance_id, worker_id)
        current_time = time.time()

        # Check if already syncing with different sync_id
        existing_sync = self._get_sync_info(instance_id, worker_id)
        if existing_sync is not None and existing_sync.state == FullSyncState.SYNCING:
            if existing_sync.sync_id != sync_id:
                logger.warning(
                    "Worker %s already syncing with different sync_id: "
                    "existing=%s, new=%s",
                    report_id,
                    existing_sync.sync_id,
                    sync_id,
                )
                return False

        new_sync_info = WorkerSyncInfo(
            sync_id=sync_id,
            state=FullSyncState.SYNCING,
            start_time=current_time,
            expected_total_keys=total_keys,
            expected_batch_count=batch_count,
            last_activity_time=current_time,
        )
        if not self._set_sync_info(instance_id, worker_id, new_sync_info):
            logger.warning(
                "Failed to start sync for worker %s: worker not found", report_id
            )
            return False

        logger.info(
            "Started full sync for worker %s: sync_id=%s, "
            "expected_keys=%d, expected_batches=%d",
            report_id,
            sync_id,
            total_keys,
            batch_count,
        )
        return True

    def receive_batch(
        self,
        instance_id: str,
        worker_id: int,
        sync_id: str,
        batch_id: int,
        keys_count: int,
    ) -> bool:
        """
        Record receipt of a sync batch.

        Returns:
            True if batch was recorded, False if invalid
        """
        report_id = (instance_id, worker_id)
        sync_info = self._get_sync_info(instance_id, worker_id)

        if sync_info is None:
            logger.warning(
                "Received batch for unknown sync session: worker=%s, sync_id=%s",
                report_id,
                sync_id,
            )
            return False

        if sync_info.sync_id != sync_id:
            logger.warning(
                "Sync ID mismatch: expected=%s, received=%s",
                sync_info.sync_id,
                sync_id,
            )
            return False

        if sync_info.state != FullSyncState.SYNCING:
            logger.warning(
                "Received batch for non-syncing worker: worker=%s, state=%s",
                report_id,
                sync_info.state,
            )
            return False

        sync_info.received_batches.add(batch_id)
        sync_info.received_keys_count += keys_count
        sync_info.last_activity_time = time.time()

        logger.debug(
            "Received batch %d for worker %s: keys=%d, total_received=%d",
            batch_id,
            report_id,
            keys_count,
            sync_info.received_keys_count,
        )
        return True

    def complete_sync(
        self,
        instance_id: str,
        worker_id: int,
        sync_id: str,
        actual_total_keys: int,
    ) -> bool:
        """
        Mark sync as completed for a worker.

        Returns:
            True if completion was successful, False otherwise
        """
        report_id = (instance_id, worker_id)
        sync_info = self._get_sync_info(instance_id, worker_id)

        if sync_info is None:
            logger.warning(
                "Received sync end for unknown session: worker=%s, sync_id=%s",
                report_id,
                sync_id,
            )
            return False

        if sync_info.sync_id != sync_id:
            logger.warning(
                "Sync ID mismatch on completion: expected=%s, received=%s",
                sync_info.sync_id,
                sync_id,
            )
            return False

        # Verify key count
        if sync_info.received_keys_count != actual_total_keys:
            logger.warning(
                "Key count mismatch on completion: received=%d, reported=%d",
                sync_info.received_keys_count,
                actual_total_keys,
            )
            # Still mark as completed but log the discrepancy

        sync_info.state = FullSyncState.COMPLETED
        sync_info.last_activity_time = time.time()

        logger.info(
            "Completed full sync for worker %s: sync_id=%s, "
            "received_keys=%d, batches=%d",
            report_id,
            sync_id,
            sync_info.received_keys_count,
            len(sync_info.received_batches),
        )
        return True

    def mark_failed(self, instance_id: str, worker_id: int, reason: str) -> None:
        """Mark a worker's sync as failed"""
        report_id = (instance_id, worker_id)
        sync_info = self._get_sync_info(instance_id, worker_id)

        if sync_info is not None:
            sync_info.state = FullSyncState.FAILED
            logger.warning(
                "Marked sync as failed for worker %s: reason=%s", report_id, reason
            )

    # Only for testing
    def check_sync_timeout(self) -> None:
        """
        Check for sync timeouts and mark failed workers.

        This should be called periodically (e.g., in health check loop).
        """
        current_time = time.time()
        for instance_id, worker_node in self._get_all_workers_cached():
            sync_info = worker_node.sync_info
            if sync_info is not None and sync_info.state == FullSyncState.SYNCING:
                if current_time - sync_info.last_activity_time > self.sync_timeout_s:
                    self.mark_failed(
                        instance_id,
                        worker_node.worker_id,
                        f"timeout after {self.sync_timeout_s}s",
                    )

    def get_global_progress(self) -> float:
        """
        Get the global sync progress.

        Returns:
            Progress as a float between 0.0 and 1.0

        Note: Uses cached worker list for Prometheus metrics efficiency.

        Progress calculation:
        - Denominator: total number of all workers
        - Numerator: workers that are ready to serve (COMPLETED or no sync info needed)
        - Workers in SYNCING or FAILED state are NOT considered ready
        """
        all_workers = self._get_all_workers_cached()
        if not all_workers:
            return 0.0

        total = len(all_workers)
        ready_count = sum(
            1
            for _, worker_node in all_workers
            if worker_node.sync_info is None
            or worker_node.sync_info.state == FullSyncState.COMPLETED
        )

        # Progress = ready / total
        return ready_count / total

    def get_completed_count(self) -> int:
        """Get count of workers that have completed sync.

        Note: Uses cached worker list for Prometheus metrics efficiency.
        """
        return sum(
            1
            for _, worker_node in self._get_all_workers_cached()
            if worker_node.sync_info is not None
            and worker_node.sync_info.state == FullSyncState.COMPLETED
        )

    def get_syncing_count(self) -> int:
        """Get count of workers currently syncing.

        Note: Uses cached worker list for Prometheus metrics efficiency.
        """
        return sum(
            1
            for _, worker_node in self._get_all_workers_cached()
            if worker_node.sync_info is not None
            and worker_node.sync_info.state == FullSyncState.SYNCING
        )

    def can_exit_freeze(self, progress: Optional[float] = None) -> bool:
        """
        Check if the completion threshold is reached and freeze mode can be exited.

        Args:
            progress: Pre-computed global progress. If None, will be computed.

        Returns:
            True if enough workers have completed sync
        """
        if progress is None:
            progress = self.get_global_progress()
        can_exit = progress >= self.completion_threshold

        if can_exit and self._need_full_sync_all:
            # Once threshold is reached, disable the global full sync flag
            logger.info(
                "Full sync completion threshold reached (%.1f%%), "
                "disabling need_full_sync_all",
                progress * 100,
            )
            self._need_full_sync_all = False

        return can_exit

    def get_total_missing_batches_count(self) -> int:
        """
        Get total count of missing batches across all syncing workers.

        Returns:
            Total number of missing batches

        Note: Uses cached worker list for Prometheus metrics efficiency.
        """
        total = 0
        for instance_id, worker_node in self._get_all_workers_cached():
            sync_info = worker_node.sync_info
            if sync_info is not None and sync_info.state == FullSyncState.SYNCING:
                expected_batches = set(range(sync_info.expected_batch_count))
                missing = expected_batches - sync_info.received_batches
                total += len(missing)
        return total

    def get_missing_batches(
        self, instance_id: str, worker_id: int, sync_id: str
    ) -> list[int]:
        """
        Get list of missing batch IDs that need to be resent.

        Args:
            instance_id: The instance ID
            worker_id: The worker ID
            sync_id: The sync session ID

        Returns:
            List of missing batch IDs, empty if sync is complete or invalid
        """
        sync_info = self._get_sync_info(instance_id, worker_id)

        if sync_info is None:
            return []

        # Check sync_id matches
        if sync_info.sync_id != sync_id:
            return []

        # If already completed, no missing batches
        if sync_info.state == FullSyncState.COMPLETED:
            return []

        # If not syncing, no missing batches
        if sync_info.state != FullSyncState.SYNCING:
            return []

        # Calculate missing batches
        expected_batches = set(range(sync_info.expected_batch_count))
        missing = expected_batches - sync_info.received_batches

        return sorted(missing)

    def get_sync_status(
        self, instance_id: str, worker_id: int, sync_id: str
    ) -> Tuple[bool, float, bool, list[int]]:
        """
        Get sync status for a specific worker.

        Returns:
            Tuple of (is_complete, global_progress, can_exit_freeze, missing_batches)
        """
        sync_info = self._get_sync_info(instance_id, worker_id)

        is_complete = (
            sync_info is not None
            and sync_info.sync_id == sync_id
            and sync_info.state == FullSyncState.COMPLETED
        )
        global_progress = self.get_global_progress()
        can_exit = self.can_exit_freeze(global_progress)
        missing_batches = self.get_missing_batches(instance_id, worker_id, sync_id)

        return is_complete, global_progress, can_exit, missing_batches