"""Diagnostics support for Tuya."""

from typing import Any

from tuya_device_handlers.helpers.diagnostics import customer_device_as_dict
from tuya_sharing import CustomerDevice

from homeassistant.components.diagnostics import REDACTED
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.device_registry import DeviceEntry

from .const import DOMAIN, DPCode
from .coordinator import TuyaConfigEntry

_REDACTED_DPCODES = {
    DPCode.ALARM_MESSAGE,
    DPCode.ALARM_MSG,
    DPCode.DOORBELL_PIC,
    DPCode.MOVEMENT_DETECT_PIC,
}


async def async_get_config_entry_diagnostics(
    hass: HomeAssistant, entry: TuyaConfigEntry
) -> dict[str, Any]:
    """Return diagnostics for a config entry."""
    return _async_get_diagnostics(hass, entry)


async def async_get_device_diagnostics(
    hass: HomeAssistant, entry: TuyaConfigEntry, device: DeviceEntry
) -> dict[str, Any]:
    """Return diagnostics for a device entry."""
    return _async_get_diagnostics(hass, entry, device)


@callback
def _async_get_diagnostics(
    hass: HomeAssistant,
    entry: TuyaConfigEntry,
    device: DeviceEntry | None = None,
) -> dict[str, Any]:
    """Return diagnostics for a config entry."""
    manager = entry.runtime_data.manager

    mqtt_connected = None
    if manager.mq.client:
        mqtt_connected = manager.mq.client.is_connected()

    data = {
        "endpoint": manager.customer_api.endpoint,
        "terminal_id": manager.terminal_id,
        "mqtt_connected": mqtt_connected,
        "disabled_by": entry.disabled_by,
        "disabled_polling": entry.pref_disable_polling,
    }

    if device:
        tuya_device_id = next(iter(device.identifiers))[1]
        data |= _async_device_as_dict(hass, entry, manager.device_map[tuya_device_id])
    else:
        data.update(
            devices=[
                _async_device_as_dict(hass, entry, device)
                for device in manager.device_map.values()
            ]
        )

    return data


@callback
def _async_device_as_dict(
    hass: HomeAssistant, entry: TuyaConfigEntry, device: CustomerDevice
) -> dict[str, Any]:
    """Represent a Tuya device as a dictionary."""

    # Base device information
    data = customer_device_as_dict(device)

    # Redact sensitive information.
    for key in data["status"]:
        if key in _REDACTED_DPCODES:
            data["status"][key] = REDACTED

    # Gather information how this Tuya device is represented in Home Assistant
    device_registry = dr.async_get(hass)
    entity_registry = er.async_get(hass)
    hass_device = device_registry.async_get_device_by_identifier(
        (DOMAIN, device.id), entry.entry_id
    )
    if hass_device:
        data["home_assistant"] = {
            "name": hass_device.name,
            "name_by_user": hass_device.name_by_user,
            "disabled": hass_device.disabled,
            "disabled_by": hass_device.disabled_by,
            "entities": [],
        }

        hass_entities = er.async_entries_for_device(
            entity_registry,
            device_id=hass_device.id,
            include_disabled_entities=True,
        )

        for entity_entry in hass_entities:
            state = hass.states.get(entity_entry.entity_id)
            state_dict: dict[str, Any] | None = None
            if state:
                state_dict = dict(state.as_dict())

                # Redact the `entity_picture` attribute as it contains a token.
                if "entity_picture" in state_dict["attributes"]:
                    state_dict["attributes"] = {
                        **state_dict["attributes"],
                        "entity_picture": REDACTED,
                    }

                # The context doesn't provide useful information in this case.
                state_dict.pop("context", None)

            data["home_assistant"]["entities"].append(
                {
                    "disabled": entity_entry.disabled,
                    "disabled_by": entity_entry.disabled_by,
                    "entity_category": entity_entry.entity_category,
                    "device_class": entity_entry.device_class,
                    "original_device_class": entity_entry.original_device_class,
                    "icon": entity_entry.icon,
                    "original_icon": entity_entry.original_icon,
                    "unit_of_measurement": entity_entry.unit_of_measurement,
                    "state": state_dict,
                }
            )

    return data