from typing import List, Dict, Optional, Any
from .state_machine import FaultState, FaultSeverity
from .detectors import FaultDetector
from .recovery import RecoveryStrategy
from altrus.core.registry.module import Module, ModuleState, ModuleHealth
from altrus.core.lifecycle import SystemState, SystemStateMachine
import time


class FaultEngine:
    """
    Severity-aware fault detection, correlation, and escalation engine.
    """

    def __init__(
        self,
        registry,
        intent_engine,
        detectors,
        recovery_strategy,
        state_machine: Optional[SystemStateMachine] = None,
        max_retries=3
    ):
        self.registry = registry
        self.intent_engine = intent_engine
        self.detectors = detectors
        self.recovery_strategy = recovery_strategy
        self.state_machine = state_machine
        self.state = FaultState.NORMAL
        self.retry_count = 0
        self.max_retries = max_retries

        # 📊 Fault Tracking & Escalation
        self.fault_history: List[Dict[str, Any]] = []
        self._max_history = 100
        self._module_failure_counts: Dict[str, int] = {}
        self._global_failure_count = 0

    def _record_fault(self, detector: FaultDetector):
        """Record fault for correlation analysis"""
        fault_info = {
            "timestamp": time.time(),
            "detector": detector.__class__.__name__,
            "fault_type": str(detector),
            "severity": detector.severity().name if hasattr(detector.severity(), 'name') else str(detector.severity()),
            "details": str(detector)
        }
        self.fault_history.append(fault_info)
        if len(self.fault_history) > self._max_history:
            self.fault_history.pop(0)

    def _analyze_root_cause(self, detected: List[FaultDetector]) -> str:
        """Simple correlation-based root cause analysis"""
        if not detected:
            return "None"

        # Count detector types
        counts = {}
        for d in detected:
            name = d.__class__.__name__
            counts[name] = counts.get(name, 0) + 1

        # Most frequent detector type as primary suspect
        primary_suspect = max(counts, key=counts.get)

        # Check for cascading failures (many detectors at once)
        if len(detected) > 3:
            return f"CASCADING_FAILURE (Primary: {primary_suspect})"

        return primary_suspect

    def handle_fault(self, module_id: str, fault_type: str, metadata: Optional[Dict] = None):
        """
        Manually signal a fault (used by simulation or internal components).
        """
        severity = FaultSeverity.SOFT
        if "CRITICAL" in fault_type.upper() or "CRASH" in fault_type.upper():
            severity = FaultSeverity.CRITICAL
        elif "HARD" in fault_type.upper():
            severity = FaultSeverity.HARD

        # Record manual fault
        fault_info = {
            "timestamp": time.time(),
            "detector": "MANUAL_SIGNAL",
            "module_id": module_id,
            "fault_type": fault_type,
            "severity": severity.name,
            "details": metadata or {}
        }
        self.fault_history.append(fault_info)

        # Log to ledger if available
        if hasattr(self.intent_engine, 'ledger'):
            self.intent_engine.ledger.log_event(
                event_type="FAULT_DETECTED",
                actor_type="MODULE",
                actor_id=module_id,
                data={
                    "fault_type": fault_type,
                    "severity": severity.name,
                    "metadata": metadata
                }
            )

        # Trigger evaluation or direct handling if critical
        effective_severity = self._escalate_severity(module_id, severity)

        if effective_severity == FaultSeverity.CRITICAL:
            self.state = FaultState.FAILED
            if self.state_machine:
                self.state_machine.transition(
                    SystemState.FAILED,
                    reason=f"Critical fault (escalated): {fault_type} for {module_id}"
                )
            # Fail the module
            self.registry.update_state(module_id, ModuleState.FAILED)
            self.registry.update_health(module_id, ModuleHealth.CRITICAL)
            self.intent_engine.on_module_failed(module_id)
        elif effective_severity == FaultSeverity.HARD:
            if self.state_machine and self.state_machine.current_state == SystemState.RUNNING:
                self.state_machine.transition(
                    SystemState.DEGRADED,
                    reason=f"Hard fault (escalated): {fault_type} for {module_id}"
                )

        return self.state

    def _escalate_severity(self, module_id: Optional[str], initial_severity: FaultSeverity) -> FaultSeverity:
        """Logic for progressive escalation: SOFT -> HARD -> CRITICAL"""
        key = module_id or "GLOBAL"
        count = self._module_failure_counts.get(key, 0)

        effective = initial_severity

        # Progressive escalation based on failure count
        if count >= 2:
            effective = FaultSeverity.CRITICAL
        elif count >= 1:
            if initial_severity == FaultSeverity.SOFT:
                effective = FaultSeverity.HARD

        # Increment for next time if it's a persistent fault
        # (Note: evaluate() usually calls this once per cycle if still detected)
        self._module_failure_counts[key] = count + 1
        return effective

    def reset_recovery(self, module_id: Optional[str] = None):
        """Recovery normalization: reset counters when system recovers"""
        if module_id:
            if module_id in self._module_failure_counts:
                del self._module_failure_counts[module_id]
        else:
             self._module_failure_counts.clear()
             self._global_failure_count = 0

        self.retry_count = 0
        self.state = FaultState.NORMAL

        # If system was degraded or failed, attempt to move back to running
        if self.state_machine and self.state_machine.current_state in {SystemState.DEGRADED, SystemState.FAILED}:
            # Check if any other faults remain
            if not any(d.check() for d in self.detectors):
                self.state_machine.transition(
                    SystemState.RUNNING,
                    reason="All faults resolved. Normalizing system state."
                )

    def evaluate(self) -> FaultState:
        """
        Run a single evaluation cycle of all registered fault detectors.

        Performs root cause analysis and triggers state transitions if
        critical faults are detected.

        Returns:
            FaultState: The updated global engine state.
        """
        detected = [d for d in self.detectors if d.check()]

        # ... logic continues ...

        if not detected:
            self.state = FaultState.NORMAL
            self.retry_count = 0
            return self.state

        # Record faults for correlation
        for d in detected:
            self._record_fault(d)

        root_cause = self._analyze_root_cause(detected)

        # Determine effective severities after escalation
        escalated_faults = []
        for d in detected:
            m_id = d.module_id()
            eff_sev = self._escalate_severity(m_id, d.severity())
            escalated_faults.append((d, eff_sev))

            if eff_sev.value > d.severity().value:
                self._log_escalation(d, d.severity(), eff_sev)

        highest_sev = max(f[1].value for f in escalated_faults)
        highest_enum = FaultSeverity(highest_sev)

        if highest_enum == FaultSeverity.CRITICAL:
            self.state = FaultState.FAILED
            if self.state_machine:
                self.state_machine.transition(
                    SystemState.FAILED,
                    reason=f"Critical fault detected (Escalated): {root_cause}"
                )
            # Fail active modules
            for module in self.registry.list_modules():
                if module.state == ModuleState.ACTIVE:
                    self.registry.update_state(module.module_id, ModuleState.FAILED)
                    self.registry.update_health(module.module_id, ModuleHealth.CRITICAL)
                    self.intent_engine.on_module_failed(module.module_id)
            return self.state

        if highest_enum == FaultSeverity.HARD:
            if self.state_machine and self.state_machine.current_state == SystemState.RUNNING:
                 self.state_machine.transition(
                    SystemState.DEGRADED,
                    reason=f"Hard fault detected (Escalated): {root_cause}"
                )

        # recovery path with exponential backoff retry logic
        if self.retry_count < self.max_retries:
            self.state = FaultState.RECOVERING

            # Exponential backoff: 2^retry_count seconds
            wait_time = 2 ** self.retry_count
            print(f"Fault recovery attempt {self.retry_count + 1} starting in {wait_time}s...")

            # Log recovery attempt
            if hasattr(self.intent_engine, 'ledger'):
                self.intent_engine.ledger.log_event(
                    event_type="RECOVERY_INITIATED",
                    actor_type="SYSTEM",
                    actor_id="fault_engine",
                    data={"attempt": self.retry_count + 1, "wait_time": wait_time}
                )

            self.recovery_strategy.execute()
            self.retry_count += 1
        else:
            self.state = FaultState.FAILED

            # Escalation to DEGRADED if max restarts exceeded
            if self.state_machine and self.state_machine.current_state == SystemState.RUNNING:
                 self.state_machine.transition(
                    SystemState.DEGRADED,
                    reason=f"Recovery failed after {self.max_retries} attempts. Root cause suspect: {root_cause}"
                )

            # Transition affected modules to DEGRADED
            for d in detected:
                m_id = d.module_id()
                if m_id:
                    self.registry.update_state(m_id, ModuleState.DEGRADED)
                    self.registry.update_health(m_id, ModuleHealth.UNSTABLE)

                    if hasattr(self.intent_engine, 'ledger'):
                        self.intent_engine.ledger.log_event(
                            event_type="MODULE_DEGRADED",
                            actor_type="MODULE",
                            actor_id=m_id,
                            data={"reason": "MAX_RESTARTS_EXCEEDED", "root_cause": root_cause}
                        )

        return self.state

    def _log_escalation(self, detector: FaultDetector, from_sev: FaultSeverity, to_sev: FaultSeverity):
        print(f"⚠️ Escalating fault severity from {from_sev.name} to {to_sev.name} for {detector}")
        event_data = {
            "detector": str(detector),
            "severity_from": from_sev.name,
            "severity_to": to_sev.name,
            "details": f"Repeat failure of {from_sev.name} fault"
        }
        self.fault_history.append({
            "timestamp": time.time(),
            "detector": "ESCALATION_ENGINE",
            "fault_type": "ESCALATION",
            "severity": to_sev.name,
            "details": event_data
        })
        if hasattr(self.registry, 'ledger'):
            self.registry.ledger.log_event(
                event_type="FAULT_ESCALATED",
                actor_type="SYSTEM",
                actor_id="fault_engine",
                data=event_data
            )

