import math

from esphome import pins
import esphome.codegen as cg
from esphome.components import canbus
from esphome.components.canbus import CONF_BIT_RATE, CanbusComponent, CanSpeed
from esphome.components.esp32 import (
    VARIANT_ESP32,
    VARIANT_ESP32C3,
    VARIANT_ESP32C5,
    VARIANT_ESP32C6,
    VARIANT_ESP32C61,
    VARIANT_ESP32H2,
    VARIANT_ESP32P4,
    VARIANT_ESP32S2,
    VARIANT_ESP32S3,
    get_esp32_variant,
    include_builtin_idf_component,
)
import esphome.config_validation as cv
from esphome.const import (
    CONF_ID,
    CONF_MODE,
    CONF_RX_PIN,
    CONF_RX_QUEUE_LEN,
    CONF_TX_PIN,
    CONF_TX_QUEUE_LEN,
)

CODEOWNERS = ["@Sympatron"]
DEPENDENCIES = ["esp32"]

CONF_TX_ENQUEUE_TIMEOUT = "tx_enqueue_timeout"

esp32_can_ns = cg.esphome_ns.namespace("esp32_can")
esp32_can = esp32_can_ns.class_("ESP32Can", CanbusComponent)

# Mode options - consistent with MCP2515 component
CanMode = esp32_can_ns.enum("CanMode")
CAN_MODES = {
    "NORMAL": CanMode.CAN_MODE_NORMAL,
    "LISTENONLY": CanMode.CAN_MODE_LISTEN_ONLY,
}

# Currently the driver only supports a subset of the bit rates defined in canbus
# The supported bit rates differ between ESP32 variants.
# See ESP-IDF Programming Guide --> API Reference --> Two-Wire Automotive Interface (TWAI)

CAN_SPEEDS_ESP32 = {
    "25KBPS": CanSpeed.CAN_25KBPS,
    "50KBPS": CanSpeed.CAN_50KBPS,
    "100KBPS": CanSpeed.CAN_100KBPS,
    "125KBPS": CanSpeed.CAN_125KBPS,
    "250KBPS": CanSpeed.CAN_250KBPS,
    "500KBPS": CanSpeed.CAN_500KBPS,
    "800KBPS": CanSpeed.CAN_800KBPS,
    "1000KBPS": CanSpeed.CAN_1000KBPS,
}

CAN_SPEEDS_ESP32_S2 = {
    "1KBPS": CanSpeed.CAN_1KBPS,
    "5KBPS": CanSpeed.CAN_5KBPS,
    "10KBPS": CanSpeed.CAN_10KBPS,
    "12K5BPS": CanSpeed.CAN_12K5BPS,
    "16KBPS": CanSpeed.CAN_16KBPS,
    "20KBPS": CanSpeed.CAN_20KBPS,
    **CAN_SPEEDS_ESP32,
}

CAN_SPEEDS_ESP32_S3 = {**CAN_SPEEDS_ESP32_S2}
CAN_SPEEDS_ESP32_C3 = {**CAN_SPEEDS_ESP32_S2}
CAN_SPEEDS_ESP32_C5 = {**CAN_SPEEDS_ESP32_S2}
CAN_SPEEDS_ESP32_C6 = {**CAN_SPEEDS_ESP32_S2}
CAN_SPEEDS_ESP32_C61 = {**CAN_SPEEDS_ESP32_S2}
CAN_SPEEDS_ESP32_H2 = {**CAN_SPEEDS_ESP32_S2}
CAN_SPEEDS_ESP32_P4 = {**CAN_SPEEDS_ESP32_S2}

CAN_SPEEDS = {
    VARIANT_ESP32: CAN_SPEEDS_ESP32,
    VARIANT_ESP32C3: CAN_SPEEDS_ESP32_C3,
    VARIANT_ESP32C5: CAN_SPEEDS_ESP32_C5,
    VARIANT_ESP32C6: CAN_SPEEDS_ESP32_C6,
    VARIANT_ESP32C61: CAN_SPEEDS_ESP32_C61,
    VARIANT_ESP32H2: CAN_SPEEDS_ESP32_H2,
    VARIANT_ESP32P4: CAN_SPEEDS_ESP32_P4,
    VARIANT_ESP32S2: CAN_SPEEDS_ESP32_S2,
    VARIANT_ESP32S3: CAN_SPEEDS_ESP32_S3,
}


def validate_bit_rate(value):
    variant = get_esp32_variant()
    if variant not in CAN_SPEEDS:
        raise cv.Invalid(f"{variant} is not supported by component {esp32_can_ns}")
    value = value.upper()
    if value not in CAN_SPEEDS[variant]:
        raise cv.Invalid(f"Bit rate {value} is not supported on {variant}")
    return cv.enum(CAN_SPEEDS[variant])(value)


CONFIG_SCHEMA = canbus.CANBUS_SCHEMA.extend(
    {
        cv.GenerateID(): cv.declare_id(esp32_can),
        cv.Optional(CONF_BIT_RATE, default="125KBPS"): validate_bit_rate,
        cv.Required(CONF_RX_PIN): pins.internal_gpio_input_pin_number,
        cv.Required(CONF_TX_PIN): pins.internal_gpio_output_pin_number,
        cv.Optional(CONF_MODE, default="NORMAL"): cv.enum(CAN_MODES, upper=True),
        cv.Optional(CONF_RX_QUEUE_LEN): cv.uint32_t,
        cv.Optional(CONF_TX_QUEUE_LEN): cv.uint32_t,
        cv.Optional(CONF_TX_ENQUEUE_TIMEOUT): cv.positive_time_period_milliseconds,
    }
)


def get_default_tx_enqueue_timeout(bit_rate):
    bit_rate_numeric = canbus.get_rate(bit_rate)
    bits_per_packet = 140  # ~max CAN message length
    ms_per_packet = bits_per_packet / bit_rate_numeric * 1000
    return int(
        max(min(math.ceil(10 * ms_per_packet), 1000), 1)
    )  # ~10 packet lengths, min 1ms, max 1000ms


async def to_code(config):
    # Legacy driver component provides driver/twai.h header
    include_builtin_idf_component("driver")
    # Also enable esp_driver_twai for future migration to new API
    include_builtin_idf_component("esp_driver_twai")
    var = cg.new_Pvariable(config[CONF_ID])
    await canbus.register_canbus(var, config)

    cg.add(var.set_rx(config[CONF_RX_PIN]))
    cg.add(var.set_tx(config[CONF_TX_PIN]))
    cg.add(var.set_mode(config[CONF_MODE]))
    if (rx_queue_len := config.get(CONF_RX_QUEUE_LEN)) is not None:
        cg.add(var.set_rx_queue_len(rx_queue_len))
    if (tx_queue_len := config.get(CONF_TX_QUEUE_LEN)) is not None:
        cg.add(var.set_tx_queue_len(tx_queue_len))

    if CONF_TX_ENQUEUE_TIMEOUT in config:
        tx_enqueue_timeout_ms = config[CONF_TX_ENQUEUE_TIMEOUT].total_milliseconds
    else:
        tx_enqueue_timeout_ms = get_default_tx_enqueue_timeout(config[CONF_BIT_RATE])

    cg.add(var.set_tx_enqueue_timeout_ms(tx_enqueue_timeout_ms))
