oon_policy.py

"""OON policy switches for UniFi Network Rules integration."""

from __future__ import annotations

import logging
from typing import Any

from homeassistant.core import callback
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.entity import generate_entity_id

from ..const import SWITCH_DELAYED_VERIFICATION_SLEEP_SECONDS
from ..coordinator import UnifiRuleUpdateCoordinator
from ..helpers.rule import (
    extract_descriptive_name,
    get_child_entity_id,
    get_child_entity_name,
    get_child_unique_id,
    get_object_id,
    get_rule_id,
)
from .base import UnifiRuleSwitch

LOGGER = logging.getLogger(__name__)


class UnifiOONPolicySwitch(UnifiRuleSwitch):
    """Switch to enable/disable a UniFi Object-Oriented Network policy."""

    def __init__(
        self,
        coordinator: UnifiRuleUpdateCoordinator,
        rule_data: Any,
        rule_type: str,
        entry_id: str = None,
    ) -> None:
        """Initialize OON policy switch."""
        super().__init__(coordinator, rule_data, rule_type, entry_id)
        # Set icon for OON policies
        self._attr_icon = "mdi:shield-network"


class UnifiOONPolicyKillSwitch(UnifiRuleSwitch):
    """Switch to enable/disable kill switch for a UniFi Object-Oriented Network policy."""

    def __init__(
        self,
        coordinator: UnifiRuleUpdateCoordinator,
        rule_data: Any,
        rule_type: str,
        entry_id: str = None,
    ) -> None:
        """Initialize OON policy kill switch using super() and overriding."""
        # 1. Initialize using the base class with the PARENT rule data.
        # The base __init__ will set initial name, unique_id etc based on the parent.
        super().__init__(coordinator, rule_data, "oon_policies", entry_id)

        # 2. Determine Parent and Kill Switch IDs
        original_rule_id = get_rule_id(rule_data)  # Parent unique ID (e.g., unr_oon_xyz)
        if not original_rule_id:
            raise ValueError("KillSwitch init: Cannot determine original rule ID from rule data")
        kill_switch_id = get_child_unique_id(original_rule_id, "kill_switch")

        # 3. Override attributes for the Kill Switch
        self._rule_id = kill_switch_id  # Internal ID used by this instance
        self._attr_unique_id = kill_switch_id  # The unique ID for HA (OVERRIDE)

        # Override name based on the name generated by super()
        # Ensure super() set a name before modifying
        if self._attr_name:
            self._attr_name = get_child_entity_name(self._attr_name, "kill_switch")  # OVERRIDE
        else:
            # Fallback name if super init failed to set one
            fallback_parent_name = extract_descriptive_name(rule_data, coordinator) or original_rule_id
            self._attr_name = get_child_entity_name(fallback_parent_name, "kill_switch")

        # Override entity_id based on the object_id generated by super()
        # Note: get_object_id uses rule_data (parent), which is correct base
        parent_object_id = get_object_id(rule_data, "oon_policies")
        kill_switch_object_id = get_child_entity_id(parent_object_id, "kill_switch")
        self.entity_id = generate_entity_id(
            "switch.{}",
            kill_switch_object_id,
            hass=coordinator.hass,  # OVERRIDE
        )

        # 4. Initialize kill switch state specifically
        # (Optimistic state handling is managed by base class, but we set initial value)
        self._optimistic_state = None
        self._optimistic_timestamp = 0
        route = getattr(rule_data, "route", {})
        if isinstance(route, dict) and "kill_switch" in route:
            actual_state = route.get("kill_switch", False)
            self._optimistic_state = actual_state  # Start optimistic state matching actual
            self._optimistic_timestamp = 0
            LOGGER.debug(
                "KillSwitch %s: Initialized specific state to %s from parent rule data", self.unique_id, actual_state
            )
        else:
            LOGGER.warning("KillSwitch %s: Initialized without specific state from rule data.", self.unique_id)

        # 5. Linking Information (Parent ID needed for lookups)
        self._linked_parent_id = original_rule_id  # Store parent's unique_id
        self._linked_child_ids = set()  # Kill switches have no children

        # Set icon for kill switch
        self._attr_icon = "mdi:shield-off"

        LOGGER.debug(
            "Finished KillSwitch __init__ for unique_id=%s, entity_id=%s", self._attr_unique_id, self.entity_id
        )

    @callback
    def _handle_coordinator_update(self) -> None:
        """Handle updated data from the coordinator.

        This method delegates to the base class implementation.
        The base class handles all the necessary logic:
        - Checking coordinator data validity
        - Finding the parent rule using our specialized _get_current_rule method
        - Handling optimistic state management
        - Triggering removal if the parent rule is missing
        """
        LOGGER.debug("KillSwitch(%s): Delegating coordinator update to base class", self.entity_id or self.unique_id)
        super()._handle_coordinator_update()

    def _get_actual_state_from_rule(self, rule: Any) -> bool | None:
        """Helper to get the actual kill switch state from the PARENT rule object."""
        if rule is None:
            LOGGER.debug(
                "KillSwitch(%s): Cannot get state, parent rule object is None.", self.entity_id or self.unique_id
            )
            return None

        # Check if the rule has the route.kill_switch attribute
        route = getattr(rule, "route", {})
        if isinstance(route, dict):
            state = route.get("kill_switch")
            if state is not None:
                return bool(state)

        LOGGER.warning(
            "KillSwitch(%s): Cannot determine actual state from parent rule object %s (type: %s)",
            self.entity_id or self.unique_id,
            getattr(rule, "id", "N/A"),
            type(rule).__name__,
        )
        return None  # Return None if state cannot be determined

    def _get_current_rule(self) -> Any | None:
        """Get the current rule from the coordinator data.

        Special implementation for KillSwitch that finds the PARENT rule in oon_policies.
        Kill switches don't have their own rule objects - they are settings on parent objects.
        """
        try:
            # 1. Validate the kill switch ID format
            kill_switch_id = self._rule_id  # This should be a kill switch ID like "unr_oon_abc_kill_switch"
            if not kill_switch_id or not kill_switch_id.endswith("_kill_switch"):
                LOGGER.warning("KillSwitch(%s): Invalid kill switch ID format: %s", self.entity_id, kill_switch_id)
                return None

            # 2. Get the parent unique ID from our stored property
            parent_unique_id = self._linked_parent_id
            if not parent_unique_id:
                LOGGER.error("KillSwitch(%s): Cannot find parent rule - _linked_parent_id is not set!", self.entity_id)
                return None

            # 3. Look for the parent rule in oon_policies collection
            # The parent's rule type is always oon_policies for kill switches
            parent_rule_type = "oon_policies"

            # Basic validations - ensure coordinator data exists
            if not self.coordinator or not self.coordinator.data or parent_rule_type not in self.coordinator.data:
                return None

            # Get the oon_policies collection
            oon_policies = self.coordinator.data[parent_rule_type]

            # Search for the parent rule by its unique ID
            parent_rule = None
            for rule in oon_policies:
                current_parent_unique_id = get_rule_id(rule)
                if current_parent_unique_id == parent_unique_id:
                    parent_rule = rule
                    break  # Found it

            # Return the parent rule (or None if not found)
            return parent_rule

        except Exception as err:
            LOGGER.error("KillSwitch(%s): Error finding parent rule: %s", self.entity_id, err)
            return None

    @property
    def is_on(self) -> bool:
        """Return true if kill switch is enabled."""
        # Use optimistic state if set
        if self._optimistic_state is not None:
            return self._optimistic_state

        current_rule = self._get_current_rule()
        if current_rule is None:
            return False

        # Check if the rule has the route.kill_switch attribute
        route = getattr(current_rule, "route", {})
        if isinstance(route, dict):
            return route.get("kill_switch", False)

        # As a last resort, default to False
        LOGGER.warning("Kill switch %s cannot determine state from rule %s", self._rule_id, type(current_rule).__name__)
        return False

    @property
    def available(self) -> bool:
        """Return if entity is available."""
        LOGGER.debug("KillSwitch(%s): Evaluating available property...", self.entity_id or self.unique_id)

        # Start with coordinator status
        coord_ok = self.coordinator.last_update_success
        if not coord_ok:
            LOGGER.debug("KillSwitch(%s): Unavailable due to coordinator failure.", self.entity_id or self.unique_id)
            return False

        # Now check if the parent rule exists in the coordinator's data
        parent_rule = self._get_current_rule()
        is_available = parent_rule is not None

        # Log the availability status, especially if it's False
        if not is_available:
            LOGGER.debug(
                "KillSwitch(%s): Determined unavailable because parent rule lookup failed.",
                self.entity_id or self.unique_id,
            )
        else:
            LOGGER.debug("KillSwitch(%s): Determined available (parent rule found).", self.entity_id or self.unique_id)

        return is_available

    async def _async_toggle_rule(self, enable: bool) -> None:
        """Toggle the kill switch setting."""
        rule = self._get_current_rule()
        if rule is None:
            raise HomeAssistantError(f"Cannot find rule with ID: {self._rule_id}")

        LOGGER.debug("%s kill switch for %s", "Enabling" if enable else "Disabling", self.name)

        # Store current state for optimistic updates
        self._optimistic_state = enable
        self._optimistic_timestamp = 0
        self.async_write_ha_state()

        # Track the operation in coordinator for proper state management
        if hasattr(self.coordinator, "_pending_operations"):
            # Use a special key for kill switch operations to avoid conflicts with parent
            kill_switch_operation_id = f"{self._rule_id}_kill_switch"
            self.coordinator._pending_operations[kill_switch_operation_id] = enable

        # Register the operation with the coordinator to prevent redundant refreshes.
        change_type = "enabled" if enable else "disabled"
        entity_id = self.entity_id or f"switch.{self._rule_id}_kill_switch"
        self.coordinator.register_ha_initiated_operation(self._rule_id, entity_id, change_type)

        # Queue the toggle operation
        try:
            # Get the toggle function from the API client
            # We need to update the policy's route.kill_switch property
            from ..models.oon_policy import OONPolicy

            policy_dict = rule.to_api_dict()
            policy_dict["route"]["kill_switch"] = enable

            # Use update_oon_policy method
            toggle_func = self.coordinator.api.update_oon_policy
            updated_policy = OONPolicy(policy_dict)  # Create new policy instance with updated data

            # Queue the operation using the coordinator's queue method
            future = await self.coordinator.api.queue_api_operation(toggle_func, updated_policy)

            async def handle_operation_complete(future):
                """Handle the completion of the toggle operation."""
                try:
                    result = future.result()
                    if result:
                        LOGGER.debug("Successfully toggled kill switch for %s", self.name)
                        # Smart Verification Task
                        import asyncio

                        async def delayed_verification():
                            await asyncio.sleep(SWITCH_DELAYED_VERIFICATION_SLEEP_SECONDS)
                            kill_switch_id = self._rule_id  # The unique ID for the kill switch
                            if self.coordinator.check_and_consume_ha_initiated_operation(kill_switch_id):
                                LOGGER.warning(
                                    "Delayed verification for kill switch %s failed. Forcing refresh.", kill_switch_id
                                )
                                await self.coordinator.async_request_refresh()
                            else:
                                LOGGER.debug("Delayed verification for kill switch %s confirmed.", kill_switch_id)

                        self.hass.async_create_task(delayed_verification())
                    else:
                        LOGGER.error("Failed to toggle kill switch for %s", self.name)
                        # Revert optimistic state on failure
                        self._optimistic_state = not enable
                        self._optimistic_timestamp = 0
                        self.async_write_ha_state()

                    # Clean up pending operations
                    if hasattr(self.coordinator, "_pending_operations"):
                        kill_switch_operation_id = f"{self._rule_id}_kill_switch"
                        if kill_switch_operation_id in self.coordinator._pending_operations:
                            del self.coordinator._pending_operations[kill_switch_operation_id]

                except Exception as err:
                    LOGGER.error("Error in kill switch toggle operation: %s", str(err))

                    # Revert optimistic state on error
                    self._optimistic_state = not enable
                    self._optimistic_timestamp = 0
                    self.async_write_ha_state()

            # Add the completion callback
            future.add_done_callback(lambda f: self.hass.async_create_task(handle_operation_complete(f)))

            LOGGER.debug("Successfully queued kill switch toggle operation for rule %s", self._rule_id)

        except Exception as error:
            LOGGER.exception("Failed to queue kill switch toggle operation: %s", str(error))

            # Clean up pending operations if queuing failed
            if hasattr(self.coordinator, "_pending_operations"):
                kill_switch_operation_id = f"{self._rule_id}_kill_switch"
                if kill_switch_operation_id in self.coordinator._pending_operations:
                    del self.coordinator._pending_operations[kill_switch_operation_id]

            # Revert optimistic state
            self._optimistic_state = not enable
            self._optimistic_timestamp = 0
            self.async_write_ha_state()

            raise HomeAssistantError(f"Error toggling kill switch for {self.name}: {error}") from error