"""Component providing support for RainMachine programs and zones."""

import asyncio
from collections.abc import Awaitable, Callable, Coroutine
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Concatenate, override

from regenmaschine.errors import RainMachineError
import voluptuous as vol

from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ATTR_ID, EntityCategory
from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers import config_validation as cv, entity_platform
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
from homeassistant.helpers.typing import VolDictType

from . import RainMachineConfigEntry, RainMachineData
from .const import (
    CONF_ALLOW_INACTIVE_ZONES_TO_RUN,
    CONF_DEFAULT_ZONE_RUN_TIME,
    CONF_DURATION,
    CONF_USE_APP_RUN_TIMES,
    DATA_PROGRAMS,
    DATA_PROVISION_SETTINGS,
    DATA_RESTRICTIONS_UNIVERSAL,
    DATA_ZONES,
    DEFAULT_ZONE_RUN,
)
from .entity import RainMachineEntity, RainMachineEntityDescription
from .services import async_update_programs_and_zones
from .util import RUN_STATE_MAP, key_exists

ATTR_ACTIVITY_TYPE = "activity_type"
ATTR_AREA = "area"
ATTR_CS_ON = "cs_on"
ATTR_CURRENT_CYCLE = "current_cycle"
ATTR_CYCLES = "cycles"
ATTR_DELAY = "delay"
ATTR_DELAY_ON = "delay_on"
ATTR_FIELD_CAPACITY = "field_capacity"
ATTR_NEXT_RUN = "next_run"
ATTR_NO_CYCLES = "number_of_cycles"
ATTR_PRECIP_RATE = "sprinkler_head_precipitation_rate"
ATTR_RESTRICTIONS = "restrictions"
ATTR_SLOPE = "slope"
ATTR_SOAK = "soak"
ATTR_SOIL_TYPE = "soil_type"
ATTR_SPRINKLER_TYPE = "sprinkler_head_type"
ATTR_STATUS = "status"
ATTR_SUN_EXPOSURE = "sun_exposure"
ATTR_VEGETATION_TYPE = "vegetation_type"
ATTR_ZONES = "zones"
ATTR_ZONE_RUN_TIME = "zone_run_time_from_app"

DAYS = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"]

SOIL_TYPE_MAP = {
    0: "Not Set",
    1: "Clay Loam",
    2: "Silty Clay",
    3: "Clay",
    4: "Loam",
    5: "Sandy Loam",
    6: "Loamy Sand",
    7: "Sand",
    8: "Sandy Clay",
    9: "Silt Loam",
    10: "Silt",
    99: "Other",
}

SLOPE_TYPE_MAP = {
    0: "Not Set",
    1: "Flat",
    2: "Moderate",
    3: "High",
    4: "Very High",
    99: "Other",
}

SPRINKLER_TYPE_MAP = {
    0: "Not Set",
    1: "Popup Spray",
    2: "Rotors Low Rate",
    3: "Surface Drip",
    4: "Bubblers Drip",
    5: "Rotors High Rate",
    99: "Other",
}

SUN_EXPOSURE_MAP = {0: "Not Set", 1: "Full Sun", 2: "Partial Shade", 3: "Full Shade"}

VEGETATION_MAP = {
    0: "Not Set",
    1: "Not Set",
    2: "Cool Season Grass",
    3: "Fruit Trees",
    4: "Flowers",
    5: "Vegetables",
    6: "Citrus",
    7: "Bushes",
    9: "Drought Tolerant Plants",
    10: "Warm Season Grass",
    11: "Trees",
    99: "Other",
}


def raise_on_request_error[_T: RainMachineBaseSwitch, **_P](
    func: Callable[Concatenate[_T, _P], Awaitable[None]],
) -> Callable[Concatenate[_T, _P], Coroutine[Any, Any, None]]:
    """Define a decorator to raise on a request error."""

    async def decorator(self: _T, *args: _P.args, **kwargs: _P.kwargs) -> None:
        """Decorate."""
        try:
            await func(self, *args, **kwargs)
        except RainMachineError as err:
            raise HomeAssistantError(
                f"Error while executing {func.__name__}: {err}",
            ) from err

    return decorator


@dataclass(frozen=True, kw_only=True)
class RainMachineSwitchDescription(
    SwitchEntityDescription, RainMachineEntityDescription
):
    """Describe a RainMachine switch."""


@dataclass(frozen=True, kw_only=True)
class RainMachineActivitySwitchDescription(RainMachineSwitchDescription):
    """Describe a RainMachine activity (program/zone) switch."""

    kind: str
    uid: int


@dataclass(frozen=True, kw_only=True)
class RainMachineRestrictionSwitchDescription(RainMachineSwitchDescription):
    """Describe a RainMachine restriction switch."""

    data_key: str


TYPE_RESTRICTIONS_FREEZE_PROTECT_ENABLED = "freeze_protect_enabled"
TYPE_RESTRICTIONS_HOT_DAYS_EXTRA_WATERING = "hot_days_extra_watering"

RESTRICTIONS_SWITCH_DESCRIPTIONS = (
    RainMachineRestrictionSwitchDescription(
        key=TYPE_RESTRICTIONS_FREEZE_PROTECT_ENABLED,
        translation_key=TYPE_RESTRICTIONS_FREEZE_PROTECT_ENABLED,
        icon="mdi:snowflake-alert",
        api_category=DATA_RESTRICTIONS_UNIVERSAL,
        data_key="freezeProtectEnabled",
    ),
    RainMachineRestrictionSwitchDescription(
        key=TYPE_RESTRICTIONS_HOT_DAYS_EXTRA_WATERING,
        translation_key=TYPE_RESTRICTIONS_HOT_DAYS_EXTRA_WATERING,
        icon="mdi:heat-wave",
        api_category=DATA_RESTRICTIONS_UNIVERSAL,
        data_key="hotDaysExtraWatering",
    ),
)


async def async_setup_entry(
    hass: HomeAssistant,
    entry: RainMachineConfigEntry,
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Set up RainMachine switches based on a config entry."""
    platform = entity_platform.async_get_current_platform()

    services: tuple[tuple[str, VolDictType | None, str], ...] = (
        ("start_program", None, "async_start_program"),
        (
            "start_zone",
            {
                vol.Optional(
                    CONF_DEFAULT_ZONE_RUN_TIME, default=DEFAULT_ZONE_RUN
                ): cv.positive_int
            },
            "async_start_zone",
        ),
        ("stop_program", None, "async_stop_program"),
        ("stop_zone", None, "async_stop_zone"),
    )
    for service_name, schema, method in services:
        platform.async_register_entity_service(service_name, schema, method)

    data = entry.runtime_data
    entities: list[RainMachineBaseSwitch] = []

    for kind, api_category, switch_class, switch_enabled_class in (
        ("program", DATA_PROGRAMS, RainMachineProgram, RainMachineProgramEnabled),
        ("zone", DATA_ZONES, RainMachineZone, RainMachineZoneEnabled),
    ):
        coordinator = data.coordinators[api_category]
        for uid, activity in coordinator.data.items():
            name = activity["name"].capitalize()

            # Add a switch to start/stop the program or zone:
            entities.append(
                switch_class(
                    entry,
                    data,
                    RainMachineActivitySwitchDescription(
                        key=f"{kind}_{uid}",
                        name=name,
                        api_category=api_category,
                        kind=kind,
                        uid=uid,
                    ),
                )
            )

            # Add a switch to enabled/disable the program or zone:
            entities.append(
                switch_enabled_class(
                    entry,
                    data,
                    RainMachineActivitySwitchDescription(
                        key=f"{kind}_{uid}_enabled",
                        name=f"{name} enabled",
                        api_category=api_category,
                        kind=kind,
                        uid=uid,
                    ),
                )
            )

    # Add switches to control restrictions:
    for description in RESTRICTIONS_SWITCH_DESCRIPTIONS:
        coordinator = data.coordinators[description.api_category]
        if not key_exists(coordinator.data, description.data_key):
            continue
        entities.append(RainMachineRestrictionSwitch(entry, data, description))

    async_add_entities(entities)


class RainMachineBaseSwitch(RainMachineEntity, SwitchEntity):
    """Define a base RainMachine switch."""

    entity_description: RainMachineSwitchDescription

    def __init__(
        self,
        entry: ConfigEntry,
        data: RainMachineData,
        description: RainMachineSwitchDescription,
    ) -> None:
        """Initialize."""
        super().__init__(entry, data, description)

        self._attr_is_on = False
        self._entry = entry

    @callback
    def _update_activities(self) -> None:
        """Update all activity data."""
        self.hass.async_create_task(
            async_update_programs_and_zones(self.hass, self._entry)
        )

    async def async_start_program(self) -> None:
        """Execute the start_program entity service."""
        raise NotImplementedError("Service not implemented for this entity")

    async def async_start_zone(self, *, zone_run_time: int) -> None:
        """Execute the start_zone entity service."""
        raise NotImplementedError("Service not implemented for this entity")

    async def async_stop_program(self) -> None:
        """Execute the stop_program entity service."""
        raise NotImplementedError("Service not implemented for this entity")

    async def async_stop_zone(self) -> None:
        """Execute the stop_zone entity service."""
        raise NotImplementedError("Service not implemented for this entity")


class RainMachineActivitySwitch(RainMachineBaseSwitch):
    """Define a RainMachine switch to start/stop an activity (program or zone)."""

    _attr_icon = "mdi:water"
    entity_description: RainMachineActivitySwitchDescription

    def __init__(
        self,
        entry: ConfigEntry,
        data: RainMachineData,
        description: RainMachineSwitchDescription,
    ) -> None:
        """Initialize."""
        super().__init__(entry, data, description)

        self._attr_extra_state_attributes[ATTR_ACTIVITY_TYPE] = (
            self.entity_description.kind
        )

    @override
    async def async_turn_off(self, **kwargs: Any) -> None:
        """Turn the switch off.

        The only way this could occur is if someone rapidly turns a disabled activity
        off right after turning it on.
        """
        if (
            not self._entry.options[CONF_ALLOW_INACTIVE_ZONES_TO_RUN]
            and not self.coordinator.data[self.entity_description.uid]["active"]
        ):
            raise HomeAssistantError(
                f"Cannot turn off an inactive program/zone: {self.name}"
            )

        await self.async_turn_off_when_active(**kwargs)

    @raise_on_request_error
    async def async_turn_off_when_active(self, **kwargs: Any) -> None:
        """Turn the switch off when its associated activity is active."""
        raise NotImplementedError

    @override
    async def async_turn_on(self, **kwargs: Any) -> None:
        """Turn the switch on."""
        if (
            not self._entry.options[CONF_ALLOW_INACTIVE_ZONES_TO_RUN]
            and not self.coordinator.data[self.entity_description.uid]["active"]
        ):
            self._attr_is_on = False
            self.async_write_ha_state()
            raise HomeAssistantError(
                f"Cannot turn on an inactive program/zone: {self.name}"
            )

        await self.async_turn_on_when_active(**kwargs)

    @raise_on_request_error
    async def async_turn_on_when_active(self, **kwargs: Any) -> None:
        """Turn the switch on when its associated activity is active."""
        raise NotImplementedError


class RainMachineEnabledSwitch(RainMachineBaseSwitch):
    """Define a RainMachine switch to enable/disable an activity (program or zone)."""

    _attr_entity_category = EntityCategory.CONFIG
    _attr_icon = "mdi:cog"
    entity_description: RainMachineActivitySwitchDescription

    def __init__(
        self,
        entry: ConfigEntry,
        data: RainMachineData,
        description: RainMachineSwitchDescription,
    ) -> None:
        """Initialize."""
        super().__init__(entry, data, description)

        self._attr_extra_state_attributes[ATTR_ACTIVITY_TYPE] = (
            self.entity_description.kind
        )

    @callback
    @override
    def update_from_latest_data(self) -> None:
        """Update the entity when new data is received."""
        self._attr_is_on = self.coordinator.data[self.entity_description.uid]["active"]


class RainMachineProgram(RainMachineActivitySwitch):
    """Define a RainMachine program."""

    @override
    async def async_start_program(self) -> None:
        """Start the program."""
        await self.async_turn_on()

    @override
    async def async_stop_program(self) -> None:
        """Stop the program."""
        await self.async_turn_off()

    @raise_on_request_error
    @override
    async def async_turn_off_when_active(self, **kwargs: Any) -> None:
        """Turn the switch off when its associated activity is active."""
        await self._data.controller.programs.stop(self.entity_description.uid)
        self._update_activities()

    @raise_on_request_error
    @override
    async def async_turn_on_when_active(self, **kwargs: Any) -> None:
        """Turn the switch on when its associated activity is active."""
        await self._data.controller.programs.start(self.entity_description.uid)
        self._update_activities()

    @callback
    @override
    def update_from_latest_data(self) -> None:
        """Update the entity when new data is received."""
        data = self.coordinator.data[self.entity_description.uid]

        self._attr_is_on = bool(data["status"])

        next_run: str | None
        if data.get("nextRun") is None:
            next_run = None
        else:
            next_run = datetime.strptime(
                f"{data['nextRun']} {data['startTime']}",
                "%Y-%m-%d %H:%M",
            ).isoformat()

        self._attr_extra_state_attributes.update(
            {
                ATTR_ID: self.entity_description.uid,
                ATTR_NEXT_RUN: next_run,
                ATTR_SOAK: data.get("soak"),
                ATTR_STATUS: RUN_STATE_MAP[data["status"]],
                ATTR_ZONES: [z for z in data["wateringTimes"] if z["active"]],
            }
        )


class RainMachineProgramEnabled(RainMachineEnabledSwitch):
    """Define a switch to enable/disable a RainMachine program."""

    @raise_on_request_error
    @override
    async def async_turn_off(self, **kwargs: Any) -> None:
        """Disable the program."""
        tasks = [
            self._data.controller.programs.stop(self.entity_description.uid),
            self._data.controller.programs.disable(self.entity_description.uid),
        ]
        await asyncio.gather(*tasks)
        self._update_activities()

    @raise_on_request_error
    @override
    async def async_turn_on(self, **kwargs: Any) -> None:
        """Enable the program."""
        await self._data.controller.programs.enable(self.entity_description.uid)
        self._update_activities()


class RainMachineRestrictionSwitch(RainMachineBaseSwitch):
    """Define a RainMachine restriction setting."""

    _attr_entity_category = EntityCategory.CONFIG
    entity_description: RainMachineRestrictionSwitchDescription

    @raise_on_request_error
    @override
    async def async_turn_off(self, **kwargs: Any) -> None:
        """Disable the restriction."""
        await self._data.controller.restrictions.set_universal(
            {self.entity_description.data_key: False}
        )
        self._attr_is_on = False
        self.async_write_ha_state()

    @raise_on_request_error
    @override
    async def async_turn_on(self, **kwargs: Any) -> None:
        """Enable the restriction."""
        await self._data.controller.restrictions.set_universal(
            {self.entity_description.data_key: True}
        )
        self._attr_is_on = True
        self.async_write_ha_state()

    @callback
    @override
    def update_from_latest_data(self) -> None:
        """Update the entity when new data is received."""
        self._attr_is_on = self.coordinator.data[self.entity_description.data_key]


class RainMachineZone(RainMachineActivitySwitch):
    """Define a RainMachine zone."""

    @override
    async def async_start_zone(self, *, zone_run_time: int) -> None:
        """Start a particular zone for a certain amount of time."""
        await self.async_turn_on(duration=zone_run_time)

    @override
    async def async_stop_zone(self) -> None:
        """Stop a zone."""
        await self.async_turn_off()

    @raise_on_request_error
    @override
    async def async_turn_off_when_active(self, **kwargs: Any) -> None:
        """Turn the switch off when its associated activity is active."""
        await self._data.controller.zones.stop(self.entity_description.uid)
        self._update_activities()

    @raise_on_request_error
    @override
    async def async_turn_on_when_active(self, **kwargs: Any) -> None:
        """Turn the switch on when its associated activity is active."""
        # 1. Use duration parameter if provided from service call
        duration = kwargs.get(CONF_DURATION)
        if not duration:
            if (
                self._entry.options[CONF_USE_APP_RUN_TIMES]
                and ATTR_ZONE_RUN_TIME in self._attr_extra_state_attributes
            ):
                # 2. Use app's zone-specific default, if enabled and available
                duration = self._attr_extra_state_attributes[ATTR_ZONE_RUN_TIME]
            else:
                # 3. Fall back to global zone default duration
                duration = self._entry.options[CONF_DEFAULT_ZONE_RUN_TIME]
        await self._data.controller.zones.start(
            self.entity_description.uid,
            duration,
        )
        self._update_activities()

    @callback
    @override
    def update_from_latest_data(self) -> None:
        """Update the entity when new data is received."""
        data = self.coordinator.data[self.entity_description.uid]

        self._attr_is_on = bool(data["state"])

        attrs = {
            ATTR_CURRENT_CYCLE: data["cycle"],
            ATTR_ID: data["uid"],
            ATTR_NO_CYCLES: data["noOfCycles"],
            ATTR_RESTRICTIONS: data["restriction"],
            ATTR_SLOPE: SLOPE_TYPE_MAP.get(data["slope"], 99),
            ATTR_SOIL_TYPE: SOIL_TYPE_MAP.get(data["soil"], 99),
            ATTR_SPRINKLER_TYPE: SPRINKLER_TYPE_MAP.get(data["group_id"], 99),
            ATTR_STATUS: RUN_STATE_MAP[data["state"]],
            ATTR_SUN_EXPOSURE: SUN_EXPOSURE_MAP.get(data.get("sun")),
            ATTR_VEGETATION_TYPE: VEGETATION_MAP.get(data["type"], 99),
        }

        if "waterSense" in data:
            if "area" in data["waterSense"]:
                attrs[ATTR_AREA] = round(data["waterSense"]["area"], 2)
            if "fieldCapacity" in data["waterSense"]:
                attrs[ATTR_FIELD_CAPACITY] = round(
                    data["waterSense"]["fieldCapacity"], 2
                )
            if "precipitationRate" in data["waterSense"]:
                attrs[ATTR_PRECIP_RATE] = round(
                    data["waterSense"]["precipitationRate"], 2
                )

        if self._entry.options[CONF_USE_APP_RUN_TIMES]:
            provision_data = self._data.coordinators[DATA_PROVISION_SETTINGS].data
            if zone_durations := provision_data.get("system", {}).get("zoneDuration"):
                attrs[ATTR_ZONE_RUN_TIME] = zone_durations[
                    list(self.coordinator.data).index(self.entity_description.uid)
                ]

        self._attr_extra_state_attributes.update(attrs)


class RainMachineZoneEnabled(RainMachineEnabledSwitch):
    """Define a switch to enable/disable a RainMachine zone."""

    @raise_on_request_error
    @override
    async def async_turn_off(self, **kwargs: Any) -> None:
        """Disable the zone."""
        tasks = [
            self._data.controller.zones.stop(self.entity_description.uid),
            self._data.controller.zones.disable(self.entity_description.uid),
        ]
        await asyncio.gather(*tasks)
        self._update_activities()

    @raise_on_request_error
    @override
    async def async_turn_on(self, **kwargs: Any) -> None:
        """Enable the zone."""
        await self._data.controller.zones.enable(self.entity_description.uid)
        self._update_activities()
