import logging
import math

from esphome import automation
from esphome.automation import Condition
import esphome.codegen as cg
from esphome.components.const import CONF_USE_PSRAM
from esphome.components.esp32 import (
    add_idf_sdkconfig_option,
    const,
    get_esp32_variant,
    only_on_variant,
)
from esphome.components.network import (
    has_high_performance_networking,
    ip_address_literal,
)
from esphome.components.psram import is_guaranteed as psram_is_guaranteed
from esphome.config_helpers import filter_source_files_from_platform
import esphome.config_validation as cv
from esphome.const import (
    CONF_AP,
    CONF_BSSID,
    CONF_CERTIFICATE,
    CONF_CERTIFICATE_AUTHORITY,
    CONF_CHANNEL,
    CONF_DNS1,
    CONF_DNS2,
    CONF_DOMAIN,
    CONF_EAP,
    CONF_ENABLE_BTM,
    CONF_ENABLE_ON_BOOT,
    CONF_ENABLE_RRM,
    CONF_FAST_CONNECT,
    CONF_GATEWAY,
    CONF_HIDDEN,
    CONF_ID,
    CONF_IDENTITY,
    CONF_KEY,
    CONF_MANUAL_IP,
    CONF_NETWORKS,
    CONF_ON_CONNECT,
    CONF_ON_DISCONNECT,
    CONF_ON_ERROR,
    CONF_OUTPUT_POWER,
    CONF_PASSWORD,
    CONF_POWER_SAVE_MODE,
    CONF_PRIORITY,
    CONF_REBOOT_TIMEOUT,
    CONF_SSID,
    CONF_STATIC_IP,
    CONF_SUBNET,
    CONF_TIMEOUT,
    CONF_TTLS_PHASE_2,
    CONF_USE_ADDRESS,
    CONF_USERNAME,
    Platform,
    PlatformFramework,
)
from esphome.core import CORE, CoroPriority, HexInt, coroutine_with_priority
import esphome.final_validate as fv
from esphome.types import ConfigType

from . import wpa2_eap

_LOGGER = logging.getLogger(__name__)

AUTO_LOAD = ["network"]

_LOGGER = logging.getLogger(__name__)

NO_WIFI_VARIANTS = [const.VARIANT_ESP32H2, const.VARIANT_ESP32P4]
CONF_SAVE = "save"
CONF_BAND_MODE = "band_mode"
CONF_MIN_AUTH_MODE = "min_auth_mode"
CONF_POST_CONNECT_ROAMING = "post_connect_roaming"

# Maximum number of WiFi networks that can be configured
# Limited to 127 because selected_sta_index_ is int8_t in C++
MAX_WIFI_NETWORKS = 127

# Default AP timeout - allows sufficient time to try all BSSIDs during initial connection
# After AP starts, WiFi scanning is skipped to avoid disrupting the AP, so we only
# get best-effort connection attempts. Longer timeout ensures we exhaust all options
# before falling back to AP mode. Aligned with improv wifi_timeout default.
DEFAULT_AP_TIMEOUT = "90s"

wifi_ns = cg.esphome_ns.namespace("wifi")
EAPAuth = wifi_ns.struct("EAPAuth")
ManualIP = wifi_ns.struct("ManualIP")
WiFiComponent = wifi_ns.class_("WiFiComponent", cg.Component)
WiFiAP = wifi_ns.struct("WiFiAP")

WiFiPowerSaveMode = wifi_ns.enum("WiFiPowerSaveMode")
WIFI_POWER_SAVE_MODES = {
    "NONE": WiFiPowerSaveMode.WIFI_POWER_SAVE_NONE,
    "LIGHT": WiFiPowerSaveMode.WIFI_POWER_SAVE_LIGHT,
    "HIGH": WiFiPowerSaveMode.WIFI_POWER_SAVE_HIGH,
}

WiFiBandMode = cg.global_ns.enum("wifi_band_mode_t")
WIFI_BAND_MODES = {
    "AUTO": WiFiBandMode.WIFI_BAND_MODE_AUTO,
    "2.4GHZ": WiFiBandMode.WIFI_BAND_MODE_2G_ONLY,
    "5GHZ": WiFiBandMode.WIFI_BAND_MODE_5G_ONLY,
}

WifiMinAuthMode = wifi_ns.enum("WifiMinAuthMode")
WIFI_MIN_AUTH_MODES = {
    "WPA": WifiMinAuthMode.WIFI_MIN_AUTH_MODE_WPA,
    "WPA2": WifiMinAuthMode.WIFI_MIN_AUTH_MODE_WPA2,
    "WPA3": WifiMinAuthMode.WIFI_MIN_AUTH_MODE_WPA3,
}
VALIDATE_WIFI_MIN_AUTH_MODE = cv.enum(WIFI_MIN_AUTH_MODES, upper=True)
WiFiConnectedCondition = wifi_ns.class_("WiFiConnectedCondition", Condition)
WiFiEnabledCondition = wifi_ns.class_("WiFiEnabledCondition", Condition)
WiFiAPActiveCondition = wifi_ns.class_("WiFiAPActiveCondition", Condition)
WiFiEnableAction = wifi_ns.class_("WiFiEnableAction", automation.Action)
WiFiDisableAction = wifi_ns.class_("WiFiDisableAction", automation.Action)
WiFiConfigureAction = wifi_ns.class_(
    "WiFiConfigureAction", automation.Action, cg.Component
)


def validate_password(value):
    value = cv.string_strict(value)
    if not value:
        return value
    if len(value) < 8:
        raise cv.Invalid("WPA password must be at least 8 characters long")
    if len(value) > 64:
        raise cv.Invalid("WPA password must be at most 64 characters long")
    return value


def validate_channel(value):
    value = cv.positive_int(value)
    if value < 1:
        raise cv.Invalid("Minimum WiFi channel is 1")
    if value > 14:
        raise cv.Invalid("Maximum WiFi channel is 14")
    return value


AP_MANUAL_IP_SCHEMA = cv.Schema(
    {
        cv.Required(CONF_STATIC_IP): cv.ipv4address,
        cv.Required(CONF_GATEWAY): cv.ipv4address,
        cv.Required(CONF_SUBNET): cv.ipv4address,
    }
)

STA_MANUAL_IP_SCHEMA = AP_MANUAL_IP_SCHEMA.extend(
    {
        cv.Optional(CONF_DNS1, default="0.0.0.0"): cv.ipv4address,
        cv.Optional(CONF_DNS2, default="0.0.0.0"): cv.ipv4address,
    }
)

TTLS_PHASE_2 = {
    "pap": cg.global_ns.ESP_EAP_TTLS_PHASE2_PAP,
    "chap": cg.global_ns.ESP_EAP_TTLS_PHASE2_CHAP,
    "mschap": cg.global_ns.ESP_EAP_TTLS_PHASE2_MSCHAP,
    "mschapv2": cg.global_ns.ESP_EAP_TTLS_PHASE2_MSCHAPV2,
    "eap": cg.global_ns.ESP_EAP_TTLS_PHASE2_EAP,
}

EAP_AUTH_SCHEMA = cv.All(
    cv.only_on([Platform.ESP32, Platform.ESP8266]),
    cv.Schema(
        {
            cv.Optional(CONF_IDENTITY): cv.string_strict,
            cv.Optional(CONF_USERNAME): cv.string_strict,
            cv.Optional(CONF_PASSWORD): cv.string_strict,
            cv.Optional(CONF_CERTIFICATE_AUTHORITY): wpa2_eap.validate_certificate,
            cv.SplitDefault(CONF_TTLS_PHASE_2, esp32="mschapv2"): cv.All(
                cv.enum(TTLS_PHASE_2), cv.only_on_esp32
            ),
            cv.Inclusive(
                CONF_CERTIFICATE, "certificate_and_key"
            ): wpa2_eap.validate_certificate,
            # Only validate as file first because we need the password to load it
            # Actual validation happens in validate_eap.
            cv.Inclusive(CONF_KEY, "certificate_and_key"): cv.file_,
        }
    ),
    wpa2_eap.validate_eap,
    cv.has_at_least_one_key(CONF_IDENTITY, CONF_CERTIFICATE),
)

WIFI_NETWORK_BASE = cv.Schema(
    {
        cv.GenerateID(): cv.declare_id(WiFiAP),
        cv.Optional(CONF_SSID): cv.ssid,
        cv.Optional(CONF_PASSWORD): validate_password,
        cv.Optional(CONF_CHANNEL): validate_channel,
        cv.Optional(CONF_MANUAL_IP): STA_MANUAL_IP_SCHEMA,
    }
)

CONF_AP_TIMEOUT = "ap_timeout"
WIFI_NETWORK_AP = WIFI_NETWORK_BASE.extend(
    {
        cv.Optional(
            CONF_AP_TIMEOUT, default=DEFAULT_AP_TIMEOUT
        ): cv.positive_time_period_milliseconds,
    }
)


def wifi_network_ap(value):
    if value is None:
        value = {}
    config = WIFI_NETWORK_AP(value)
    if CONF_MANUAL_IP in config and CORE.is_rp2040:
        raise cv.Invalid(
            "Manual AP IP configuration is not supported on RP2040. "
            "The AP uses the default IP 192.168.4.1"
        )
    return config


WIFI_NETWORK_STA = WIFI_NETWORK_BASE.extend(
    {
        cv.Optional(CONF_BSSID): cv.mac_address,
        cv.Optional(CONF_HIDDEN): cv.boolean,
        cv.Optional(CONF_PRIORITY, default=0): cv.int_range(min=-128, max=127),
        cv.Optional(CONF_EAP): EAP_AUTH_SCHEMA,
    }
)


def validate_variant(_):
    if CORE.is_esp32:
        variant = get_esp32_variant()
        if variant in NO_WIFI_VARIANTS and "esp32_hosted" not in fv.full_config.get():
            raise cv.Invalid(f"WiFi requires component esp32_hosted on {variant}")
    if CORE.is_rp2040:
        from esphome.components.rp2040 import board_has_wifi, get_board

        if not board_has_wifi():
            raise cv.Invalid(
                f"Board '{get_board()}' does not have WiFi support (no CYW43 wireless chip). "
                f"Use a WiFi-capable board like 'rpipicow' or 'rpipico2w'."
            )


def _apply_min_auth_mode_default(config):
    """Apply platform-specific default for min_auth_mode and warn ESP8266 users."""
    # Only apply defaults for platforms that support min_auth_mode
    if CONF_MIN_AUTH_MODE not in config and (CORE.is_esp8266 or CORE.is_esp32):
        if CORE.is_esp8266:
            _LOGGER.warning(
                "The minimum WiFi authentication mode (wifi -> min_auth_mode) is not set. "
                "This controls the weakest encryption your device will accept when connecting to WiFi. "
                "Currently defaults to WPA (less secure), but will change to WPA2 (more secure) in 2026.6.0. "
                "WPA uses TKIP encryption which has known security vulnerabilities and should be avoided. "
                "WPA2 uses AES encryption which is significantly more secure. "
                "To silence this warning, explicitly set min_auth_mode under 'wifi:'. "
                "If your router supports WPA2 or WPA3, set 'min_auth_mode: WPA2'. "
                "If your router only supports WPA, set 'min_auth_mode: WPA'."
            )
            config[CONF_MIN_AUTH_MODE] = VALIDATE_WIFI_MIN_AUTH_MODE("WPA")
        elif CORE.is_esp32:
            config[CONF_MIN_AUTH_MODE] = VALIDATE_WIFI_MIN_AUTH_MODE("WPA2")
    return config


def final_validate(config):
    has_sta = bool(config.get(CONF_NETWORKS, True))
    has_ap = CONF_AP in config
    full_config = fv.full_config.get()
    has_improv = "esp32_improv" in full_config
    has_improv_serial = "improv_serial" in full_config
    has_captive_portal = "captive_portal" in full_config
    has_web_server = "web_server" in full_config
    if not (has_sta or has_ap or has_improv or has_improv_serial):
        raise cv.Invalid(
            "Please specify at least an SSID or an Access Point to create."
        )
    if has_ap and not has_captive_portal and not has_web_server:
        _LOGGER.warning(
            "WiFi AP is configured but neither captive_portal nor web_server is enabled. "
            "The AP will not be usable for configuration or monitoring. "
            "Add 'captive_portal:' or 'web_server:' to your configuration."
        )


def _consume_wifi_sockets(config: ConfigType) -> ConfigType:
    """Register UDP PCBs used internally by lwIP for DHCP and DNS.

    Only needed on LibreTiny where we directly set MEMP_NUM_UDP_PCB (the raw
    PCB pool shared by both application sockets and lwIP internals like DHCP/DNS).
    On ESP32, CONFIG_LWIP_MAX_SOCKETS only controls the POSIX socket layer —
    DHCP/DNS use raw udp_new() which bypasses it entirely.
    """
    if not (CORE.is_bk72xx or CORE.is_rtl87xx or CORE.is_ln882x):
        return config
    from esphome.components import socket

    # lwIP allocates UDP PCBs for DHCP client and DNS resolver internally.
    # These are not application sockets but consume MEMP_NUM_UDP_PCB pool entries.
    socket.consume_sockets(2, "wifi.lwip_internal", socket.SocketType.UDP)(config)
    return config


FINAL_VALIDATE_SCHEMA = cv.All(
    final_validate,
    validate_variant,
    _consume_wifi_sockets,
)


def _validate(config):
    if CONF_PASSWORD in config and CONF_SSID not in config:
        raise cv.Invalid("Cannot have WiFi password without SSID!")

    if CONF_SSID in config:
        # Automatically move single network to 'networks' section
        config = config.copy()
        network = {CONF_SSID: config.pop(CONF_SSID)}
        if CONF_PASSWORD in config:
            network[CONF_PASSWORD] = config.pop(CONF_PASSWORD)
        if CONF_EAP in config:
            network[CONF_EAP] = config.pop(CONF_EAP)
        if CONF_NETWORKS in config:
            # In testing mode, merged component tests may have both ssid and networks
            # Just use the networks list and ignore the single ssid
            if not CORE.testing_mode:
                raise cv.Invalid(
                    "You cannot use the 'ssid:' option together with 'networks:'. Please "
                    "copy your network into the 'networks:' key"
                )
        else:
            config[CONF_NETWORKS] = cv.ensure_list(WIFI_NETWORK_STA)(network)

    if (CONF_NETWORKS not in config) and (CONF_AP not in config):
        config = config.copy()
        config[CONF_NETWORKS] = []

    if CONF_USE_ADDRESS not in config:
        use_address = CORE.name + config[CONF_DOMAIN]
        if CONF_MANUAL_IP in config:
            use_address = str(config[CONF_MANUAL_IP][CONF_STATIC_IP])
        elif CONF_NETWORKS in config:
            ips = {
                str(net[CONF_MANUAL_IP][CONF_STATIC_IP])
                for net in config[CONF_NETWORKS]
                if CONF_MANUAL_IP in net
            }
            if len(ips) > 1:
                raise cv.Invalid(
                    "Must specify use_address when using multiple static IP addresses."
                )
            if len(ips) == 1:
                use_address = next(iter(ips))

        config[CONF_USE_ADDRESS] = use_address

    return config


CONF_PASSIVE_SCAN = "passive_scan"
CONFIG_SCHEMA = cv.All(
    cv.Schema(
        {
            cv.GenerateID(): cv.declare_id(WiFiComponent),
            cv.Optional(CONF_NETWORKS): cv.All(
                cv.ensure_list(WIFI_NETWORK_STA), cv.Length(max=MAX_WIFI_NETWORKS)
            ),
            cv.Optional(CONF_SSID): cv.ssid,
            cv.Optional(CONF_PASSWORD): validate_password,
            cv.Optional(CONF_MANUAL_IP): STA_MANUAL_IP_SCHEMA,
            cv.Optional(CONF_EAP): EAP_AUTH_SCHEMA,
            cv.Optional(CONF_AP): wifi_network_ap,
            cv.Optional(CONF_DOMAIN, default=".local"): cv.domain_name,
            cv.Optional(
                CONF_REBOOT_TIMEOUT, default="15min"
            ): cv.positive_time_period_milliseconds,
            cv.SplitDefault(
                CONF_POWER_SAVE_MODE,
                esp8266="none",
                esp32="light",
                rp2040="light",
                bk72xx="none",
                rtl87xx="none",
                ln882x="light",
            ): cv.enum(WIFI_POWER_SAVE_MODES, upper=True),
            cv.Optional(CONF_FAST_CONNECT, default=False): cv.boolean,
            cv.Optional(CONF_USE_ADDRESS): cv.string_strict,
            cv.Optional(CONF_MIN_AUTH_MODE): cv.All(
                VALIDATE_WIFI_MIN_AUTH_MODE,
                cv.only_on([Platform.ESP32, Platform.ESP8266]),
            ),
            cv.SplitDefault(CONF_OUTPUT_POWER, esp8266=20.0): cv.All(
                cv.decibel, cv.float_range(min=8.5, max=20.5)
            ),
            cv.SplitDefault(CONF_ENABLE_BTM, esp32=False): cv.All(
                cv.boolean, cv.only_on_esp32
            ),
            cv.SplitDefault(CONF_ENABLE_RRM, esp32=False): cv.All(
                cv.boolean, cv.only_on_esp32
            ),
            cv.Optional(CONF_BAND_MODE): cv.All(
                cv.enum(WIFI_BAND_MODES, upper=True),
                cv.only_on_esp32,
                only_on_variant(supported=[const.VARIANT_ESP32C5]),
            ),
            cv.Optional(CONF_PASSIVE_SCAN, default=False): cv.boolean,
            cv.Optional(CONF_ENABLE_ON_BOOT, default=True): cv.boolean,
            cv.Optional(CONF_POST_CONNECT_ROAMING, default=True): cv.boolean,
            cv.Optional(CONF_ON_CONNECT): automation.validate_automation(single=True),
            cv.Optional(CONF_ON_DISCONNECT): automation.validate_automation(
                single=True
            ),
            cv.Optional(CONF_USE_PSRAM): cv.All(
                cv.only_on_esp32, cv.requires_component("psram"), cv.boolean
            ),
        }
    ),
    _apply_min_auth_mode_default,
    _validate,
)


def eap_auth(config):
    if config is None:
        return None
    ca_cert = ""
    if CONF_CERTIFICATE_AUTHORITY in config:
        ca_cert = wpa2_eap.read_relative_config_path(config[CONF_CERTIFICATE_AUTHORITY])
    client_cert = ""
    if CONF_CERTIFICATE in config:
        client_cert = wpa2_eap.read_relative_config_path(config[CONF_CERTIFICATE])
    key = ""
    if CONF_KEY in config:
        key = wpa2_eap.read_relative_config_path(config[CONF_KEY])
    return cg.StructInitializer(
        EAPAuth,
        ("identity", config.get(CONF_IDENTITY, "")),
        ("username", config.get(CONF_USERNAME, "")),
        ("password", config.get(CONF_PASSWORD, "")),
        ("ca_cert", ca_cert),
        ("client_cert", client_cert),
        ("client_key", key),
        ("ttls_phase_2", config.get(CONF_TTLS_PHASE_2)),
    )


def safe_ip(ip):
    return ip_address_literal(ip)


def manual_ip(config):
    if config is None:
        return None
    return cg.StructInitializer(
        ManualIP,
        ("static_ip", safe_ip(config[CONF_STATIC_IP])),
        ("gateway", safe_ip(config[CONF_GATEWAY])),
        ("subnet", safe_ip(config[CONF_SUBNET])),
        ("dns1", safe_ip(config.get(CONF_DNS1))),
        ("dns2", safe_ip(config.get(CONF_DNS2))),
    )


def wifi_network(config, ap, static_ip):
    if CONF_SSID in config:
        cg.add(ap.set_ssid(config[CONF_SSID]))
    if CONF_PASSWORD in config:
        cg.add(ap.set_password(config[CONF_PASSWORD]))
    if CONF_EAP in config:
        cg.add(ap.set_eap(eap_auth(config[CONF_EAP])))
        cg.add_define("USE_WIFI_WPA2_EAP")
    if CONF_BSSID in config:
        cg.add(ap.set_bssid([HexInt(i) for i in config[CONF_BSSID].parts]))
    if CONF_HIDDEN in config:
        cg.add(ap.set_hidden(config[CONF_HIDDEN]))
    if CONF_CHANNEL in config:
        cg.add(ap.set_channel(config[CONF_CHANNEL]))
    if static_ip is not None:
        cg.add(ap.set_manual_ip(manual_ip(static_ip)))
    if CONF_PRIORITY in config:
        cg.add(ap.set_priority(config[CONF_PRIORITY]))

    return ap


@coroutine_with_priority(CoroPriority.COMMUNICATION)
async def to_code(config):
    var = cg.new_Pvariable(config[CONF_ID])
    cg.add(var.set_use_address(config[CONF_USE_ADDRESS]))

    # Track if any network uses Enterprise authentication
    has_eap = False
    # Track if any network uses manual IP
    has_manual_ip = False

    # Initialize FixedVector with the count of networks
    networks = config.get(CONF_NETWORKS, [])
    if networks:
        cg.add(var.init_sta(len(networks)))

        def add_sta(ap: cg.MockObj, network: dict) -> None:
            ip_config = network.get(CONF_MANUAL_IP, config.get(CONF_MANUAL_IP))
            cg.add(var.add_sta(wifi_network(network, ap, ip_config)))

        for network in networks:
            if CONF_EAP in network:
                has_eap = True
            if network.get(CONF_MANUAL_IP) or config.get(CONF_MANUAL_IP):
                has_manual_ip = True
            cg.with_local_variable(network[CONF_ID], WiFiAP(), add_sta, network)

    if CONF_AP in config:
        conf = config[CONF_AP]
        ip_config = conf.get(CONF_MANUAL_IP)
        if ip_config:
            has_manual_ip = True
        cg.with_local_variable(
            conf[CONF_ID],
            WiFiAP(),
            lambda ap: cg.add(var.set_ap(wifi_network(conf, ap, ip_config))),
        )
        cg.add(var.set_ap_timeout(conf[CONF_AP_TIMEOUT]))
        cg.add_define("USE_WIFI_AP")
    elif CORE.is_esp32 and not CORE.using_arduino:
        add_idf_sdkconfig_option("CONFIG_ESP_WIFI_SOFTAP_SUPPORT", False)
        add_idf_sdkconfig_option("CONFIG_LWIP_DHCPS", False)

    # Disable Enterprise WiFi support if no EAP is configured
    if CORE.is_esp32:
        add_idf_sdkconfig_option("CONFIG_ESP_WIFI_ENTERPRISE_SUPPORT", has_eap)

    # Only define USE_WIFI_MANUAL_IP if any AP uses manual IP
    if has_manual_ip:
        cg.add_define("USE_WIFI_MANUAL_IP")

    cg.add(var.set_reboot_timeout(config[CONF_REBOOT_TIMEOUT]))
    cg.add(var.set_power_save_mode(config[CONF_POWER_SAVE_MODE]))
    if CONF_MIN_AUTH_MODE in config:
        cg.add(var.set_min_auth_mode(config[CONF_MIN_AUTH_MODE]))
    if config[CONF_FAST_CONNECT]:
        cg.add_define("USE_WIFI_FAST_CONNECT")
    # passive_scan defaults to false in C++ - only set if true
    if config[CONF_PASSIVE_SCAN]:
        cg.add(var.set_passive_scan(True))
    if CONF_OUTPUT_POWER in config:
        cg.add(var.set_output_power(config[CONF_OUTPUT_POWER]))
        if CORE.is_esp32:
            # Set PHY max TX power to match output_power so calibration also uses
            # reduced power. This prevents brownout during PHY init on marginal
            # power supplies, which is critical for OTA updates with rollback enabled.
            # Kconfig range is 10-20, ESPHome allows 8.5-20.5
            phy_tx_power = max(10, min(20, math.ceil(config[CONF_OUTPUT_POWER])))
            add_idf_sdkconfig_option("CONFIG_ESP_PHY_MAX_WIFI_TX_POWER", phy_tx_power)
    # enable_on_boot defaults to true in C++ - only set if false
    if not config[CONF_ENABLE_ON_BOOT]:
        cg.add(var.set_enable_on_boot(False))

    # post_connect_roaming defaults to true in C++ - disable if user disabled it
    # or if 802.11k/v is enabled (driver handles roaming natively)
    if (
        not config[CONF_POST_CONNECT_ROAMING]
        or config.get(CONF_ENABLE_BTM)
        or config.get(CONF_ENABLE_RRM)
    ):
        cg.add(var.set_post_connect_roaming(False))

    if CORE.is_esp8266:
        cg.add_library("ESP8266WiFi", None)
    elif CORE.is_rp2040:
        cg.add_library("WiFi", None)

    if CORE.is_esp32:
        if config[CONF_ENABLE_BTM] or config[CONF_ENABLE_RRM]:
            add_idf_sdkconfig_option("CONFIG_WPA_11KV_SUPPORT", True)
            cg.add_define("USE_WIFI_11KV_SUPPORT")
        if config[CONF_ENABLE_BTM]:
            cg.add(var.set_btm(config[CONF_ENABLE_BTM]))
        if config[CONF_ENABLE_RRM]:
            cg.add(var.set_rrm(config[CONF_ENABLE_RRM]))
        if CONF_BAND_MODE in config:
            cg.add(var.set_band_mode(config[CONF_BAND_MODE]))

    if config.get(CONF_USE_PSRAM):
        add_idf_sdkconfig_option("CONFIG_SPIRAM_TRY_ALLOCATE_WIFI_LWIP", True)

    # Apply high performance WiFi settings if high performance networking is enabled
    if CORE.is_esp32 and has_high_performance_networking():
        # Check if PSRAM is guaranteed (set by psram component during final validation)
        psram_guaranteed = psram_is_guaranteed()

        # Always allocate WiFi buffers in PSRAM if available
        add_idf_sdkconfig_option("CONFIG_SPIRAM_TRY_ALLOCATE_WIFI_LWIP", True)

        if psram_guaranteed:
            _LOGGER.info(
                "Applying high-performance WiFi settings (PSRAM guaranteed): 512 RX buffers, 32 TX buffers"
            )
            # PSRAM is guaranteed - use aggressive settings
            # Higher maximum values are allowed because CONFIG_LWIP_WND_SCALE is set to true in networking component
            # Based on https://github.com/espressif/esp-adf/issues/297#issuecomment-783811702

            # Large dynamic RX buffers (requires PSRAM)
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_STATIC_RX_BUFFER_NUM", 16)
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_DYNAMIC_RX_BUFFER_NUM", 512)

            # Static TX buffers for better performance
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_STATIC_TX_BUFFER", True)
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_TX_BUFFER_TYPE", 0)
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_CACHE_TX_BUFFER_NUM", 32)
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_STATIC_TX_BUFFER_NUM", 8)

            # AMPDU settings optimized for PSRAM
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_AMPDU_TX_ENABLED", True)
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_TX_BA_WIN", 16)
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_AMPDU_RX_ENABLED", True)
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_RX_BA_WIN", 32)
        else:
            _LOGGER.info(
                "Applying optimized WiFi settings: 64 RX buffers, 64 TX buffers"
            )
            # PSRAM not guaranteed - use more conservative, but still optimized settings
            # Based on https://github.com/espressif/esp-idf/blob/release/v5.4/examples/wifi/iperf/sdkconfig.defaults.esp32

            # Standard buffer counts
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_STATIC_RX_BUFFER_NUM", 16)
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_DYNAMIC_RX_BUFFER_NUM", 64)
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_DYNAMIC_TX_BUFFER_NUM", 64)

            # Standard AMPDU settings
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_AMPDU_TX_ENABLED", True)
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_TX_BA_WIN", 32)
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_AMPDU_RX_ENABLED", True)
            add_idf_sdkconfig_option("CONFIG_ESP_WIFI_RX_BA_WIN", 32)

    cg.add_define("USE_WIFI")

    # must register before OTA safe mode check
    await cg.register_component(var, config)

    await cg.past_safe_mode()

    if on_connect_config := config.get(CONF_ON_CONNECT):
        cg.add_define("USE_WIFI_CONNECT_TRIGGER")
        await automation.build_automation(
            var.get_connect_trigger(), [], on_connect_config
        )

    if on_disconnect_config := config.get(CONF_ON_DISCONNECT):
        cg.add_define("USE_WIFI_DISCONNECT_TRIGGER")
        await automation.build_automation(
            var.get_disconnect_trigger(), [], on_disconnect_config
        )

    CORE.add_job(final_step)


@automation.register_condition("wifi.connected", WiFiConnectedCondition, cv.Schema({}))
async def wifi_connected_to_code(config, condition_id, template_arg, args):
    return cg.new_Pvariable(condition_id, template_arg)


@automation.register_condition("wifi.enabled", WiFiEnabledCondition, cv.Schema({}))
async def wifi_enabled_to_code(config, condition_id, template_arg, args):
    return cg.new_Pvariable(condition_id, template_arg)


@automation.register_condition("wifi.ap_active", WiFiAPActiveCondition, cv.Schema({}))
async def wifi_ap_active_to_code(config, condition_id, template_arg, args):
    return cg.new_Pvariable(condition_id, template_arg)


@automation.register_action(
    "wifi.enable", WiFiEnableAction, cv.Schema({}), synchronous=True
)
async def wifi_enable_to_code(config, action_id, template_arg, args):
    return cg.new_Pvariable(action_id, template_arg)


@automation.register_action(
    "wifi.disable", WiFiDisableAction, cv.Schema({}), synchronous=True
)
async def wifi_disable_to_code(config, action_id, template_arg, args):
    return cg.new_Pvariable(action_id, template_arg)


KEEP_SCAN_RESULTS_KEY = "wifi_keep_scan_results"
RUNTIME_POWER_SAVE_KEY = "wifi_runtime_power_save"
# Keys for listener counts
IP_STATE_LISTENERS_KEY = "wifi_ip_state_listeners"
SCAN_RESULTS_LISTENERS_KEY = "wifi_scan_results_listeners"
CONNECT_STATE_LISTENERS_KEY = "wifi_connect_state_listeners"
POWER_SAVE_LISTENERS_KEY = "wifi_power_save_listeners"


def request_wifi_scan_results():
    """Request that WiFi scan results be kept in memory after connection.

    Components that need access to scan results after WiFi is connected should
    call this function during their code generation. This prevents the WiFi component from
    freeing scan result memory after successful connection.
    """
    CORE.data[KEEP_SCAN_RESULTS_KEY] = True


def enable_runtime_power_save_control():
    """Enable runtime WiFi power save control.

    Components that need to dynamically switch WiFi power saving on/off for latency
    performance (e.g., audio streaming, large data transfers) should call this
    function during their code generation. This enables the request_high_performance()
    and release_high_performance() APIs.

    Only supported on ESP32.
    """
    CORE.data[RUNTIME_POWER_SAVE_KEY] = True


def request_wifi_ip_state_listener() -> None:
    """Request an IP state listener slot."""
    CORE.data[IP_STATE_LISTENERS_KEY] = CORE.data.get(IP_STATE_LISTENERS_KEY, 0) + 1


def request_wifi_scan_results_listener() -> None:
    """Request a scan results listener slot."""
    CORE.data[SCAN_RESULTS_LISTENERS_KEY] = (
        CORE.data.get(SCAN_RESULTS_LISTENERS_KEY, 0) + 1
    )


def request_wifi_connect_state_listener() -> None:
    """Request a connect state listener slot."""
    CORE.data[CONNECT_STATE_LISTENERS_KEY] = (
        CORE.data.get(CONNECT_STATE_LISTENERS_KEY, 0) + 1
    )


def request_wifi_power_save_listener() -> None:
    """Request a power save listener slot."""
    CORE.data[POWER_SAVE_LISTENERS_KEY] = CORE.data.get(POWER_SAVE_LISTENERS_KEY, 0) + 1


@coroutine_with_priority(CoroPriority.FINAL)
async def final_step():
    """Final code generation step to configure optional WiFi features."""
    if CORE.data.get(KEEP_SCAN_RESULTS_KEY, False):
        cg.add(
            cg.RawExpression("wifi::global_wifi_component->set_keep_scan_results(true)")
        )
    if CORE.data.get(RUNTIME_POWER_SAVE_KEY, False):
        cg.add_define("USE_WIFI_RUNTIME_POWER_SAVE")

    # Generate listener defines - each listener type has its own #ifdef
    ip_state_count = CORE.data.get(IP_STATE_LISTENERS_KEY, 0)
    scan_results_count = CORE.data.get(SCAN_RESULTS_LISTENERS_KEY, 0)
    connect_state_count = CORE.data.get(CONNECT_STATE_LISTENERS_KEY, 0)
    power_save_count = CORE.data.get(POWER_SAVE_LISTENERS_KEY, 0)

    if ip_state_count:
        cg.add_define("USE_WIFI_IP_STATE_LISTENERS")
        cg.add_define("ESPHOME_WIFI_IP_STATE_LISTENERS", ip_state_count)
    if scan_results_count:
        cg.add_define("USE_WIFI_SCAN_RESULTS_LISTENERS")
        cg.add_define("ESPHOME_WIFI_SCAN_RESULTS_LISTENERS", scan_results_count)
    if connect_state_count:
        cg.add_define("USE_WIFI_CONNECT_STATE_LISTENERS")
        cg.add_define("ESPHOME_WIFI_CONNECT_STATE_LISTENERS", connect_state_count)
    if power_save_count:
        cg.add_define("USE_WIFI_POWER_SAVE_LISTENERS")
        cg.add_define("ESPHOME_WIFI_POWER_SAVE_LISTENERS", power_save_count)


@automation.register_action(
    "wifi.configure",
    WiFiConfigureAction,
    cv.Schema(
        {
            cv.Required(CONF_SSID): cv.templatable(cv.ssid),
            cv.Required(CONF_PASSWORD): cv.templatable(validate_password),
            cv.Optional(CONF_SAVE, default=True): cv.templatable(cv.boolean),
            cv.Optional(CONF_TIMEOUT, default="30000ms"): cv.templatable(
                cv.positive_time_period_milliseconds
            ),
            cv.Optional(CONF_ON_CONNECT): automation.validate_automation(single=True),
            cv.Optional(CONF_ON_ERROR): automation.validate_automation(single=True),
        }
    ),
    synchronous=False,
)
async def wifi_set_sta_to_code(config, action_id, template_arg, args):
    var = cg.new_Pvariable(action_id, template_arg)
    ssid = await cg.templatable(config[CONF_SSID], args, cg.std_string)
    password = await cg.templatable(config[CONF_PASSWORD], args, cg.std_string)
    save = await cg.templatable(config[CONF_SAVE], args, cg.bool_)
    timeout = await cg.templatable(config.get(CONF_TIMEOUT), args, cg.uint32)
    cg.add(var.set_ssid(ssid))
    cg.add(var.set_password(password))
    cg.add(var.set_save(save))
    cg.add(var.set_connection_timeout(timeout))
    if on_connect_config := config.get(CONF_ON_CONNECT):
        await automation.build_automation(
            var.get_connect_trigger(), [], on_connect_config
        )
    if on_error_config := config.get(CONF_ON_ERROR):
        await automation.build_automation(var.get_error_trigger(), [], on_error_config)
    await cg.register_component(var, config)
    return var


FILTER_SOURCE_FILES = filter_source_files_from_platform(
    {
        "wifi_component_esp_idf.cpp": {
            PlatformFramework.ESP32_IDF,
            PlatformFramework.ESP32_ARDUINO,
        },
        "wifi_component_esp8266.cpp": {PlatformFramework.ESP8266_ARDUINO},
        "wifi_component_libretiny.cpp": {
            PlatformFramework.BK72XX_ARDUINO,
            PlatformFramework.RTL87XX_ARDUINO,
            PlatformFramework.LN882X_ARDUINO,
        },
        "wifi_component_pico_w.cpp": {PlatformFramework.RP2040_ARDUINO},
    }
)
