__init__.py

from esphome import automation
import esphome.codegen as cg
from esphome.components import (
    binary_sensor,
    number,
    select,
    sensor,
    switch,
    text_sensor,
    uart,
)
import esphome.config_validation as cv
from esphome.const import (
    CONF_ID,
    CONF_MODE,
    CONF_TEMPERATURE,
    DEVICE_CLASS_HUMIDITY,
    DEVICE_CLASS_PROBLEM,
    DEVICE_CLASS_TEMPERATURE,
    ENTITY_CATEGORY_DIAGNOSTIC,
    ICON_CHIP,
    STATE_CLASS_MEASUREMENT,
    STATE_CLASS_TOTAL_INCREASING,
    UNIT_CELSIUS,
    UNIT_PERCENT,
)

CODEOWNERS = ["@karl"]
DEPENDENCIES = ["uart"]
AUTO_LOAD = ["binary_sensor", "number", "select", "sensor", "switch", "text_sensor"]
MULTI_CONF = True

CONF_AUTO_STOP = "auto_stop"
CONF_CURRENT_HUMIDITY = "current_humidity"
CONF_DISPLAY = "display"
CONF_MANUAL_INTENSITY = "manual_intensity"
CONF_MCU_FIRMWARE_VERSION = "mcu_firmware_version"
CONF_ON_BUTTON_EVENT = "on_button_event"
CONF_PACKET_DECODE_FAILURES = "packet_decode_failures"
CONF_POWER = "power"
CONF_TARGET_HUMIDITY = "target_humidity"
CONF_WATER_LACKING = "water_lacking"

MODE_OPTIONS = ["off", "manual", "auto"]

levoit_humidifier_ns = cg.esphome_ns.namespace("levoit_humidifier")
LevoitHumidifier = levoit_humidifier_ns.class_(
    "LevoitHumidifier", cg.PollingComponent, uart.UARTDevice
)
LevoitModeSelect = levoit_humidifier_ns.class_("LevoitModeSelect", select.Select)
LevoitPowerSwitch = levoit_humidifier_ns.class_("LevoitPowerSwitch", switch.Switch)
LevoitDisplaySwitch = levoit_humidifier_ns.class_("LevoitDisplaySwitch", switch.Switch)
LevoitAutoStopSwitch = levoit_humidifier_ns.class_(
    "LevoitAutoStopSwitch", switch.Switch
)
LevoitTargetHumidityNumber = levoit_humidifier_ns.class_(
    "LevoitTargetHumidityNumber", number.Number
)
LevoitManualIntensityNumber = levoit_humidifier_ns.class_(
    "LevoitManualIntensityNumber", number.Number
)

CONFIG_SCHEMA = (
    cv.Schema(
        {
            cv.GenerateID(): cv.declare_id(LevoitHumidifier),
            cv.Optional(CONF_CURRENT_HUMIDITY): sensor.sensor_schema(
                unit_of_measurement=UNIT_PERCENT,
                accuracy_decimals=0,
                device_class=DEVICE_CLASS_HUMIDITY,
                state_class=STATE_CLASS_MEASUREMENT,
            ),
            cv.Optional(CONF_TEMPERATURE): sensor.sensor_schema(
                unit_of_measurement=UNIT_CELSIUS,
                accuracy_decimals=0,
                device_class=DEVICE_CLASS_TEMPERATURE,
                state_class=STATE_CLASS_MEASUREMENT,
            ),
            cv.Optional(CONF_MCU_FIRMWARE_VERSION): text_sensor.text_sensor_schema(
                entity_category=ENTITY_CATEGORY_DIAGNOSTIC, icon=ICON_CHIP
            ),
            cv.Optional(CONF_WATER_LACKING): binary_sensor.binary_sensor_schema(
                device_class=DEVICE_CLASS_PROBLEM
            ),
            cv.Optional(CONF_POWER): switch.switch_schema(LevoitPowerSwitch),
            cv.Optional(CONF_DISPLAY): switch.switch_schema(LevoitDisplaySwitch),
            cv.Optional(CONF_AUTO_STOP): switch.switch_schema(LevoitAutoStopSwitch),
            cv.Optional(CONF_MODE): select.select_schema(LevoitModeSelect),
            cv.Optional(CONF_TARGET_HUMIDITY): number.number_schema(
                LevoitTargetHumidityNumber,
                unit_of_measurement=UNIT_PERCENT,
                device_class=DEVICE_CLASS_HUMIDITY,
            ),
            cv.Optional(CONF_MANUAL_INTENSITY): number.number_schema(
                LevoitManualIntensityNumber
            ),
            cv.Optional(CONF_ON_BUTTON_EVENT): automation.validate_automation({}),
            cv.Optional(CONF_PACKET_DECODE_FAILURES): sensor.sensor_schema(
                accuracy_decimals=0,
                entity_category=ENTITY_CATEGORY_DIAGNOSTIC,
                state_class=STATE_CLASS_TOTAL_INCREASING,
            ),
        }
    )
    .extend(cv.polling_component_schema("5s"))
    .extend(uart.UART_DEVICE_SCHEMA)
)

FINAL_VALIDATE_SCHEMA = uart.final_validate_device_schema(
    "levoit_humidifier", baud_rate=9600, require_tx=True, require_rx=True
)

_CALLBACK_AUTOMATIONS = (
    automation.CallbackAutomation(
        CONF_ON_BUTTON_EVENT,
        "add_on_button_event_callback",
        [(cg.uint8, "button_code")],
    ),
)


async def to_code(config):
    var = cg.new_Pvariable(config[CONF_ID])
    await cg.register_component(var, config)
    await uart.register_uart_device(var, config)
    await automation.build_callback_automations(var, config, _CALLBACK_AUTOMATIONS)

    if current_humidity_config := config.get(CONF_CURRENT_HUMIDITY):
        sens = await sensor.new_sensor(current_humidity_config)
        cg.add(var.set_current_humidity_sensor(sens))

    if temperature_config := config.get(CONF_TEMPERATURE):
        sens = await sensor.new_sensor(temperature_config)
        cg.add(var.set_temperature_sensor(sens))

    if mcu_firmware_version_config := config.get(CONF_MCU_FIRMWARE_VERSION):
        sens = await text_sensor.new_text_sensor(mcu_firmware_version_config)
        cg.add(var.set_mcu_firmware_version_text_sensor(sens))

    if water_lacking_config := config.get(CONF_WATER_LACKING):
        sens = await binary_sensor.new_binary_sensor(water_lacking_config)
        cg.add(var.set_water_lacking_binary_sensor(sens))

    if packet_decode_failures_config := config.get(CONF_PACKET_DECODE_FAILURES):
        sens = await sensor.new_sensor(packet_decode_failures_config)
        cg.add(var.set_packet_decode_failures_sensor(sens))

    if power_config := config.get(CONF_POWER):
        entity = await switch.new_switch(power_config, var)
        cg.add(var.set_power_switch(entity))

    if display_config := config.get(CONF_DISPLAY):
        entity = await switch.new_switch(display_config, var)
        cg.add(var.set_display_switch(entity))

    if auto_stop_config := config.get(CONF_AUTO_STOP):
        entity = await switch.new_switch(auto_stop_config, var)
        cg.add(var.set_auto_stop_switch(entity))

    if mode_config := config.get(CONF_MODE):
        entity = await select.new_select(mode_config, var, options=MODE_OPTIONS)
        cg.add(var.set_mode_select(entity))

    if target_humidity_config := config.get(CONF_TARGET_HUMIDITY):
        entity = await number.new_number(
            target_humidity_config, var, min_value=30, max_value=80, step=1
        )
        cg.add(var.set_target_humidity_number(entity))

    if manual_intensity_config := config.get(CONF_MANUAL_INTENSITY):
        entity = await number.new_number(
            manual_intensity_config, var, min_value=1, max_value=9, step=1
        )
        cg.add(var.set_manual_intensity_number(entity))