"""
Copyright (c) 2025 Faustino Lopez Ramos. 
For licensing information, see the LICENSE file in the project root
"""
import threading
import time
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Set

from loguru import logger  # type: ignore


class CircuitBreakerState(Enum):
    """States of the circuit breaker."""

    CLOSED = "closed"  # Normal operation
    OPEN = "open"  # Blocking requests due to failures
    HALF_OPEN = "half_open"  # Testing if system recovered


class CircuitBreaker:
    """Circuit breaker with auto-reset and half-open state.

    Prevents cascading failures by temporarily blocking operations
    after repeated failures. Automatically attempts recovery.
    """

    def __init__(
        self,
        failure_threshold: int = 3,
        timeout: timedelta = timedelta(minutes=5),
        half_open_max_calls: int = 1,
    ):
        """Initialize circuit breaker.

        Args:
            failure_threshold: Number of failures before opening circuit
            timeout: Time to wait before attempting recovery
            half_open_max_calls: Max calls allowed in half-open state
        """
        self.failure_threshold = failure_threshold
        self.timeout = timeout
        self.half_open_max_calls = half_open_max_calls

        self._lock = threading.RLock()
        self.state = CircuitBreakerState.CLOSED
        self.failure_count = 0
        self.success_count = 0
        self.last_failure_time: Optional[datetime] = None
        self.half_open_calls = 0

    def record_success(self) -> None:
        """Record successful operation."""
        with self._lock:
            self.failure_count = 0

            if self.state == CircuitBreakerState.HALF_OPEN:
                self.success_count += 1
                logger.info(
                    f"Circuit breaker success in HALF_OPEN state "
                    f"({self.success_count}/{self.half_open_max_calls})"
                )

                # If enough successes in half-open, close the circuit
                if self.success_count >= self.half_open_max_calls:
                    self._close_circuit()

            elif self.state == CircuitBreakerState.CLOSED:
                # Reset any partial failure count
                self.failure_count = 0

    def record_failure(self) -> bool:
        """Record failed operation.

        Returns:
            True if circuit should open (threshold exceeded)
        """
        with self._lock:
            self.failure_count += 1
            self.last_failure_time = datetime.now()

            if self.state == CircuitBreakerState.HALF_OPEN:
                # Failure in half-open state reopens circuit
                logger.warning("Circuit breaker failure in HALF_OPEN state - reopening circuit")
                self._open_circuit()
                return True

            elif self.state == CircuitBreakerState.CLOSED:
                if self.failure_count >= self.failure_threshold:
                    logger.error(
                        f"Circuit breaker threshold exceeded "
                        f"({self.failure_count}/{self.failure_threshold}) - opening circuit"
                    )
                    self._open_circuit()
                    return True

            return False

    def can_execute(self) -> bool:
        """Check if execution is allowed.

        ER2 FIX: Thread-safe state machine with all state changes inside lock.
        Prevents race condition where multiple threads could obtain True in HALF_OPEN state.
        CLOSED -> OPEN -> HALF_OPEN -> CLOSED

        Returns:
            True if operation can proceed
        """
        with self._lock:
            if self.state == CircuitBreakerState.CLOSED:
                return True

            elif self.state == CircuitBreakerState.OPEN:
                # Check if timeout expired
                if self.last_failure_time is None:
                    return False

                time_since_failure = datetime.now() - self.last_failure_time
                if time_since_failure >= self.timeout:
                    logger.info(
                        f"Circuit breaker timeout expired ({time_since_failure.total_seconds():.1f}s) "
                        "- transitioning to HALF_OPEN"
                    )
                    self._transition_to_half_open()
                    # CRITICAL: Increment happens inside lock to prevent race condition
                    self.half_open_calls = 1
                    return True

                return False

            elif self.state == CircuitBreakerState.HALF_OPEN:
                # Allow limited calls in half-open state
                # CRITICAL: All operations inside lock to prevent race condition
                if self.half_open_calls < self.half_open_max_calls:
                    self.half_open_calls += 1
                    return True
                return False

            return False

    def _open_circuit(self) -> None:
        """Transition to OPEN state."""
        self.state = CircuitBreakerState.OPEN
        logger.critical("Circuit breaker OPEN - blocking operations")

    def _close_circuit(self) -> None:
        """Transition to CLOSED state (normal operation)."""
        self.state = CircuitBreakerState.CLOSED
        self.failure_count = 0
        self.success_count = 0
        self.half_open_calls = 0
        logger.info("Circuit breaker CLOSED - normal operation resumed")

    def _transition_to_half_open(self) -> None:
        """Transition to HALF_OPEN state (testing recovery)."""
        self.state = CircuitBreakerState.HALF_OPEN
        self.success_count = 0
        self.half_open_calls = 0
        logger.info("Circuit breaker HALF_OPEN - testing system recovery")

    def get_state(self) -> CircuitBreakerState:
        """Get current circuit breaker state."""
        with self._lock:
            return self.state

    def reset(self) -> None:
        """Manually reset circuit breaker to CLOSED state."""
        with self._lock:
            self._close_circuit()
            logger.info("Circuit breaker manually reset")


class NodeStatus(Enum):
    """Status of a node in the pipeline execution."""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    RETRYING = "retrying"
    CANCELLED = "cancelled"


class NodeType(Enum):
    """Type of pipeline node."""

    BATCH = "batch"
    STREAMING = "streaming"
    FEATURE_STORE = "feature_store"  # Native feature store operations


@dataclass
class NodeExecutionInfo:
    """Information about a node's execution state."""

    node_name: str
    node_type: NodeType
    status: NodeStatus = NodeStatus.PENDING
    start_time: Optional[float] = None
    end_time: Optional[float] = None
    output_path: Optional[str] = None
    error: Optional[str] = None
    dependencies: List[str] = field(default_factory=list)
    dependents: Set[str] = field(default_factory=set)
    retry_count: int = 0
    max_retries: int = 3
    retry_delay: float = 5.0  # seconds
    execution_metadata: Dict[str, Any] = field(default_factory=dict)
    resources: List[Any] = field(default_factory=list)


class UnifiedPipelineState:
    """Unified state for coordinating batch and streaming pipeline execution."""

    def __init__(
        self,
        circuit_breaker_threshold: int = 3,
        circuit_breaker_timeout_minutes: int = 5,
    ):
        """Initialize the unified pipeline state."""
        self._lock = threading.RLock()
        self._nodes: Dict[str, NodeExecutionInfo] = {}

        # New circuit breaker with auto-reset
        self._circuit_breaker = CircuitBreaker(
            failure_threshold=circuit_breaker_threshold,
            timeout=timedelta(minutes=circuit_breaker_timeout_minutes),
            half_open_max_calls=2,
        )

        # Legacy attributes (for backward compatibility)
        self._circuit_breaker_threshold = circuit_breaker_threshold
        self._failure_counts: Dict[str, int] = defaultdict(int)

        self._pipeline_status: str = "initializing"
        self._batch_outputs: Dict[str, str] = {}
        self._streaming_queries: Dict[str, Any] = {}
        self._streaming_stopper: Optional[Callable[[str], bool]] = None
        self._cross_dependencies: Dict[str, Set[str]] = {}

    def register_node(
        self,
        node_name: str,
        node_type: NodeType,
        dependencies: Optional[List[str]] = None,
    ) -> None:
        """Register a node in the unified state."""
        with self._lock:
            if node_name in self._nodes:
                raise ValueError(f"Node '{node_name}' already registered")

            deps = dependencies or []
            self._nodes[node_name] = NodeExecutionInfo(
                node_name=node_name,
                node_type=node_type,
                dependencies=list(deps),
            )

            for dep in deps:
                if dep in self._nodes:
                    self._nodes[dep].dependents.add(node_name)

            for other_name, other_info in self._nodes.items():
                if other_name == node_name:
                    continue
                if node_name in other_info.dependencies:
                    self._nodes[node_name].dependents.add(other_name)

            if node_type == NodeType.STREAMING:
                batch_deps = [
                    dep
                    for dep in deps
                    if dep in self._nodes and self._nodes[dep].node_type == NodeType.BATCH
                ]
                if batch_deps:
                    self._cross_dependencies[node_name] = set(batch_deps)

            logger.debug(
                f"Registered {node_type.value} node '{node_name}' with dependencies: {dependencies}"
            )

    def start_node_execution(self, node_name: str) -> bool:
        """Start execution of a node if dependencies are ready and circuit allows."""
        with self._lock:
            if node_name not in self._nodes:
                raise ValueError(f"Node '{node_name}' not registered")

            node = self._nodes[node_name]

            # Check circuit breaker before execution
            if not self._circuit_breaker.can_execute():
                logger.warning(
                    f"Circuit breaker blocked execution of node '{node_name}' - "
                    f"state: {self._circuit_breaker.get_state().value}"
                )
                return False

            if not self._are_dependencies_ready(node_name):
                return False

            node.status = NodeStatus.RUNNING
            node.start_time = time.time()

            logger.info(f"Started execution of {node.node_type.value} node '{node_name}'")
            return True

    def complete_node_execution(
        self,
        node_name: str,
        output_path: Optional[str] = None,
        execution_metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Mark a node as completed successfully."""
        with self._lock:
            if node_name not in self._nodes:
                raise ValueError(f"Node '{node_name}' not registered")

            node = self._nodes[node_name]
            node.status = NodeStatus.COMPLETED
            node.end_time = time.time()
            node.output_path = output_path
            node.execution_metadata = execution_metadata or {}

            # Record success in circuit breaker (helps with recovery)
            self._circuit_breaker.record_success()

            if node.node_type == NodeType.BATCH and output_path:
                self._batch_outputs[node_name] = output_path

            execution_time = node.end_time - (node.start_time or node.end_time)
            logger.info(
                f"Completed {node.node_type.value} node '{node_name}' in {execution_time:.2f}s"
            )

            self._notify_dependents(node_name)

    def fail_node_execution(self, node_name: str, error: str) -> None:
        """Handle node failure with retry logic and circuit breaker."""
        with self._lock:
            if node_name not in self._nodes:
                raise ValueError(f"Node '{node_name}' not registered")

            node = self._nodes[node_name]
            self._failure_counts[node_name] += 1

            # Record failure in circuit breaker
            should_open = self._circuit_breaker.record_failure()

            if should_open:
                logger.critical(
                    f"Circuit breaker triggered for node '{node_name}' - "
                    f"state: {self._circuit_breaker.get_state().value}"
                )
                node.status = NodeStatus.FAILED
                node.error = f"Circuit breaker triggered: {error}"
                node.end_time = time.time()
                self._propagate_failure(node_name)
                return

            if node.retry_count < node.max_retries:
                node.retry_count += 1
                node.status = NodeStatus.RETRYING
                logger.warning(
                    f"Node '{node_name}' failed (attempt {node.retry_count}/"
                    f"{node.max_retries}). Retrying in {node.retry_delay}s"
                )
                threading.Timer(node.retry_delay, self._retry_node, args=[node_name]).start()
            else:
                node.status = NodeStatus.FAILED
                node.error = error
                node.end_time = time.time()
                self._propagate_failure(node_name)

    def _retry_node(self, node_name: str) -> None:
        """Retry node execution after failure."""
        with self._lock:
            if node_name not in self._nodes:
                return

            node = self._nodes[node_name]
            if node.status == NodeStatus.RETRYING:
                node.status = NodeStatus.PENDING
                logger.info(f"Retrying node '{node_name}'")

    def register_streaming_query(self, node_name: str, query: Any) -> None:
        """Register a streaming query for tracking."""
        with self._lock:
            self._streaming_queries[node_name] = query
            logger.debug(f"Registered streaming query for node '{node_name}'")

    def set_streaming_stopper(self, stopper: Callable[[str], bool]) -> None:
        """Register a stopper function to stop queries by execution_id (string)."""
        with self._lock:
            self._streaming_stopper = stopper

    def get_batch_output_path(self, node_name: str) -> Optional[str]:
        """Get the output path of a batch node."""
        with self._lock:
            return self._batch_outputs.get(node_name)

    def is_node_ready(self, node_name: str) -> bool:
        """Check if a node is ready to execute."""
        with self._lock:
            return self._are_dependencies_ready(node_name)

    def get_node_status(self, node_name: str) -> Optional[NodeStatus]:
        """Get the status of a node."""
        with self._lock:
            node = self._nodes.get(node_name)
            return node.status if node else None

    def get_ready_nodes(self) -> List[str]:
        """Get list of nodes ready for execution."""
        with self._lock:
            ready_nodes = []
            for node_name, node in self._nodes.items():
                if node.status == NodeStatus.PENDING and self._are_dependencies_ready(node_name):
                    ready_nodes.append(node_name)
            return ready_nodes

    def get_pipeline_summary(self) -> Dict[str, Any]:
        """Get summary of the pipeline state."""
        with self._lock:
            summary = {
                "total_nodes": len(self._nodes),
                "batch_nodes": len(
                    [n for n in self._nodes.values() if n.node_type == NodeType.BATCH]
                ),
                "streaming_nodes": len(
                    [n for n in self._nodes.values() if n.node_type == NodeType.STREAMING]
                ),
                "status_breakdown": {},
                "cross_dependencies": len(self._cross_dependencies),
                "pipeline_status": self._pipeline_status,
                "circuit_breaker_state": self._circuit_breaker.get_state().value,
                "circuit_breaker_failure_count": self._circuit_breaker.failure_count,
                # Legacy field for backward compatibility
                "circuit_breaker_open": self._circuit_breaker.get_state()
                == CircuitBreakerState.OPEN,
            }

            for status in NodeStatus:
                count = len([n for n in self._nodes.values() if n.status == status])
                summary["status_breakdown"][status.value] = count

            return summary

    def validate_cross_dependencies(self) -> List[str]:
        """Validate cross-dependencies between batch and streaming nodes."""
        warnings = []

        with self._lock:
            for streaming_node, batch_deps in self._cross_dependencies.items():
                for batch_dep in batch_deps:
                    batch_node = self._nodes.get(batch_dep)
                    if not batch_node:
                        warnings.append(
                            f"Streaming node '{streaming_node}' depends on "
                            f"non-existent batch node '{batch_dep}'"
                        )
                    elif batch_node.node_type != NodeType.BATCH:
                        warnings.append(
                            f"Cross-dependency error: '{streaming_node}' -> '{batch_dep}' "
                            f"but '{batch_dep}' is not a batch node"
                        )

        return warnings

    def _attempt_stop_by_handle(
        self,
        streaming_node: str,
        streaming_query: Any,
        stopped_nodes: List[str],
        failed_batch_node: str,
    ) -> None:
        try:
            streaming_query.stop()
            self._nodes[streaming_node].status = NodeStatus.CANCELLED
            stopped_nodes.append(streaming_node)
            logger.warning(
                f"Stopped streaming node '{streaming_node}' due to "
                f"failed dependency '{failed_batch_node}'"
            )
        except Exception as e:
            logger.error(f"Failed to stop streaming node '{streaming_node}': {e}")

    def _attempt_stop_by_id(
        self,
        streaming_node: str,
        query_id: str,
        stopped_nodes: List[str],
        failed_batch_node: str,
    ) -> None:
        try:
            if not self._streaming_stopper:
                logger.error(f"No stopper configured to stop streaming node '{streaming_node}'")
                return

            if self._streaming_stopper(query_id):
                self._nodes[streaming_node].status = NodeStatus.CANCELLED
                stopped_nodes.append(streaming_node)
                logger.warning(
                    f"Stopped streaming node '{streaming_node}' by id due to "
                    f"failed dependency '{failed_batch_node}'"
                )
            else:
                logger.error(f"Stopper failed to stop streaming node '{streaming_node}'")
        except Exception as e:
            logger.error(f"Error invoking stopper for '{streaming_node}': {e}")

    def stop_dependent_streaming_nodes(self, failed_batch_node: str) -> List[str]:
        """Stop streaming nodes that depend on a failed batch node."""
        stopped_nodes = []

        with self._lock:
            for streaming_node, batch_deps in self._cross_dependencies.items():
                if failed_batch_node not in batch_deps:
                    continue

                streaming_query = self._streaming_queries.get(streaming_node)
                if not streaming_query:
                    continue

                if hasattr(streaming_query, "stop"):
                    self._attempt_stop_by_handle(
                        streaming_node,
                        streaming_query,
                        stopped_nodes,
                        failed_batch_node,
                    )
                elif isinstance(streaming_query, str):
                    self._attempt_stop_by_id(
                        streaming_node,
                        streaming_query,
                        stopped_nodes,
                        failed_batch_node,
                    )

        return stopped_nodes

    def _are_dependencies_ready(self, node_name: str) -> bool:
        """Check if all dependencies of a node are completed."""
        node = self._nodes.get(node_name)
        if not node:
            return False

        for dep in node.dependencies:
            dep_node = self._nodes.get(dep)
            if not dep_node or dep_node.status != NodeStatus.COMPLETED:
                return False

        return True

    def _notify_dependents(self, completed_node: str) -> None:
        """Notify dependent nodes that a node has completed."""
        node = self._nodes.get(completed_node)
        if not node:
            return

        for dependent in node.dependents:
            if self._are_dependencies_ready(dependent):
                logger.debug(
                    f"Node '{dependent}' is now ready (dependency '{completed_node}' completed)"
                )

    def _propagate_failure(self, failed_node: str) -> None:
        """Propagate failure to dependent nodes with circuit breaker awareness."""
        circuit_state = self._circuit_breaker.get_state()
        if circuit_state == CircuitBreakerState.OPEN:
            logger.error(
                "Circuit breaker OPEN - propagating failure to all dependent nodes. "
                f"System will attempt recovery in {self._circuit_breaker.timeout.total_seconds():.0f}s"
            )

        node = self._nodes.get(failed_node)
        if node and node.node_type == NodeType.BATCH:
            self.stop_dependent_streaming_nodes(failed_node)

    def set_pipeline_status(self, status: str) -> None:
        """Set the overall pipeline status."""
        with self._lock:
            self._pipeline_status = status
            logger.debug(f"Pipeline status changed to: {status}")

    def register_resource(self, node_name: str, resource_type: str, resource: Any):
        """Register a resource for cleanup tracking"""
        with self._lock:
            if node_name not in self._nodes:
                raise ValueError(f"Node '{node_name}' not registered")

            self._nodes[node_name].resources.append((resource_type, resource))
            logger.debug(f"Registered {resource_type} resource for node '{node_name}'")

    def _stop_query_by_handle(self, node_name: str, query: Any) -> None:
        try:
            query.stop()
            logger.info(f"Stopped streaming query handle for node '{node_name}'")
        except Exception as e:
            logger.warning(f"Error stopping streaming query handle for '{node_name}': {e}")

    def _stop_query_by_id(self, node_name: str, query_id: str) -> None:
        try:
            if self._streaming_stopper:
                self._streaming_stopper(query_id)
                logger.info(f"Requested stop by id for streaming node '{node_name}'")
            else:
                logger.error(f"No stopper configured to stop streaming node '{node_name}'")
        except Exception as e:
            logger.warning(f"Error requesting stop by id for '{node_name}': {e}")

    def _stop_query(self, node_name: str, query: Any) -> None:
        if hasattr(query, "stop"):
            self._stop_query_by_handle(node_name, query)
        elif isinstance(query, str):
            self._stop_query_by_id(node_name, query)

    def _release_resource(self, node, res_type: str, resource: Any) -> None:
        try:
            if res_type == "spark_rdd" and hasattr(resource, "unpersist"):
                resource.unpersist()
            elif hasattr(resource, "close"):
                resource.close()
            logger.debug(f"Released {res_type} for node '{node.node_name}'")
        except Exception as e:
            logger.error(f"Error releasing resource for '{node.node_name}': {str(e)}")

    def cleanup(self) -> None:
        """Clean up all resources and stop active streaming queries"""
        with self._lock:
            for node_name, query in self._streaming_queries.items():
                self._stop_query(node_name, query)

            for node in self._nodes.values():
                for res_type, resource in node.resources:
                    self._release_resource(node, res_type, resource)

            self._reset_state()
            logger.debug("Pipeline state has been reset")

    def _reset_state(self):
        """Reset all state containers"""
        self._streaming_queries.clear()
        self._nodes.clear()
        self._batch_outputs.clear()
        self._cross_dependencies.clear()
        self._failure_counts.clear()
        self._pipeline_status = "initializing"
        self._circuit_open = False
