"""ESPHome display platform that dispatches to GxEPD2's panel driver classes.

model: is validated against components/gxepd2/catalog.json, generated by
tools/gxepd2_catalog/generate_catalog.py.
"""

import json
from pathlib import Path

import esphome.codegen as cg
import esphome.config_validation as cv
from esphome import pins
from esphome.components import display, spi
from esphome.components.esp32 import add_idf_component
from esphome.const import (
    CONF_BUSY_PIN,
    CONF_CS_PIN,
    CONF_DC_PIN,
    CONF_ID,
    CONF_MODEL,
    CONF_RESET_PIN,
    CONF_SPI_ID,
    Framework,
)
from esphome.core import CORE
from esphome.types import ConfigType

CODEOWNERS = ["@zib"]
DEPENDENCIES = ["spi"]

# Not in esphome/const.py: only the dual-CS "master+slave" panels need it.
CONF_CS2_PIN = "cs2_pin"

VENDOR_ROOT = Path(__file__).parent / "vendor"
CATALOG_PATH = Path(__file__).parent / "catalog.json"
CATALOG = json.loads(CATALOG_PATH.read_text())

# gxepd2::GxEPD2ColorMode enumerator name for each catalog wrapper class.
COLOR_MODES = {
    "GxEPD2_BW": "BW",
    "GxEPD2_3C": "THREE_COLOR",
    "GxEPD2_4C": "FOUR_COLOR",
    "GxEPD2_7C": "SEVEN_COLOR",
}
# USE_GXEPD2_* define gating which wrapper header gxepd2_display.h includes.
WRAPPER_DEFINES = {
    "GxEPD2_BW": "USE_GXEPD2_BW",
    "GxEPD2_3C": "USE_GXEPD2_3C",
    "GxEPD2_4C": "USE_GXEPD2_4C",
    "GxEPD2_7C": "USE_GXEPD2_7C",
}
# IT8951 host-controller panels need genuine bidirectional SPI (they read
# status/data back over MISO); arduino_compat/SPI.h's transfer() is
# write-only - see the comment there - so these can never work through it.
MODELS = {
    model_id: entry
    for model_id, entry in CATALOG.items()
    if entry["wrapper"] in COLOR_MODES and not entry["source"].startswith("it8951/")
}

gxepd2_ns = cg.esphome_ns.namespace("gxepd2")
GxEPD2Display = gxepd2_ns.class_("GxEPD2Display", display.DisplayBuffer)

# GxEPD2's plain int16_t pin ids (PinRegistry keys), assigned at codegen time.
_PIN_ID_KEY = "gxepd2_next_pin_id"


def _allocate_pin_id() -> int:
    pin_id = CORE.data.get(_PIN_ID_KEY, 0)
    CORE.data[_PIN_ID_KEY] = pin_id + 1
    return pin_id


# .cpp files this config needs, accumulated across every gxepd2 instance.
_SOURCES_KEY = "gxepd2_needed_sources"
_SHARED_SOURCE = "GxEPD2_EPD.cpp"


def _write_sources_manifest(source: str) -> None:
    needed: set[str] = CORE.data.setdefault(_SOURCES_KEY, {_SHARED_SOURCE})
    needed.add(source)
    manifest = CORE.relative_src_path("gxepd2_sources.cmake")
    manifest.parent.mkdir(parents=True, exist_ok=True)
    lines = ["set(GXEPD2_SOURCES"]
    lines += [f'    "${{CMAKE_CURRENT_SOURCE_DIR}}/src/{path}"' for path in sorted(needed)]
    lines.append(")")
    manifest.write_text("\n".join(lines) + "\n")


def validate_cs2_pin_required(config: ConfigType) -> ConfigType:
    needs_cs2 = MODELS[config[CONF_MODEL]]["ctor_args"] == 5
    has_cs2 = CONF_CS2_PIN in config
    if needs_cs2 and not has_cs2:
        raise cv.Invalid(f"'{CONF_CS2_PIN}' is required for model {config[CONF_MODEL]} (a dual-CS panel)")
    if has_cs2 and not needs_cs2:
        raise cv.Invalid(f"'{CONF_CS2_PIN}' is not used by model {config[CONF_MODEL]}")
    return config


CONFIG_SCHEMA = cv.All(
    display.FULL_DISPLAY_SCHEMA.extend(
        {
            cv.GenerateID(): cv.declare_id(GxEPD2Display),
            cv.GenerateID(CONF_SPI_ID): cv.use_id(spi.SPIComponent),
            cv.Required(CONF_MODEL): cv.one_of(*MODELS, upper=False),
            cv.Required(CONF_CS_PIN): pins.gpio_output_pin_schema,
            cv.Required(CONF_DC_PIN): pins.gpio_output_pin_schema,
            cv.Optional(CONF_RESET_PIN): pins.gpio_output_pin_schema,
            cv.Optional(CONF_BUSY_PIN): pins.gpio_input_pin_schema,
            cv.Optional(CONF_CS2_PIN): pins.gpio_output_pin_schema,
        }
    ),
    cv.only_with_framework(Framework.ESP_IDF),
    validate_cs2_pin_required,
)


async def to_code(config: ConfigType) -> None:
    entry = MODELS[config[CONF_MODEL]]
    driver_t = cg.global_ns.class_(entry["driver"])
    page_height = f"{entry['driver']}::HEIGHT / {entry['height_divisor']}"
    display_t = cg.RawExpression(f"{entry['wrapper']}<{entry['driver']}, {page_height}>")
    color_mode = cg.RawExpression(f"gxepd2::GxEPD2ColorMode::{COLOR_MODES[entry['wrapper']]}")
    dual_cs = entry["ctor_args"] == 5
    specific_type = GxEPD2Display.template(driver_t, display_t, color_mode, dual_cs)

    cg.add_define(WRAPPER_DEFINES[entry["wrapper"]])

    add_idf_component(name=VENDOR_ROOT.name, path=str(VENDOR_ROOT))
    _write_sources_manifest(entry["source"])

    cs_id, dc_id = _allocate_pin_id(), _allocate_pin_id()
    reset_id = _allocate_pin_id() if CONF_RESET_PIN in config else -1
    busy_id = _allocate_pin_id() if CONF_BUSY_PIN in config else -1
    ctor_args = [cs_id, dc_id, reset_id, busy_id]
    if entry["ctor_args"] == 5:
        cs2_id = _allocate_pin_id()
        ctor_args.append(cs2_id)

    rhs = specific_type.new(*ctor_args)
    var = cg.Pvariable(config[CONF_ID], rhs, specific_type)

    cg.add(var.set_pin(cs_id, await cg.gpio_pin_expression(config[CONF_CS_PIN])))
    cg.add(var.set_pin(dc_id, await cg.gpio_pin_expression(config[CONF_DC_PIN])))
    if reset_pin := config.get(CONF_RESET_PIN):
        cg.add(var.set_pin(reset_id, await cg.gpio_pin_expression(reset_pin)))
    if busy_pin := config.get(CONF_BUSY_PIN):
        cg.add(var.set_pin(busy_id, await cg.gpio_pin_expression(busy_pin)))
    if entry["ctor_args"] == 5:
        cg.add(var.set_pin(cs2_id, await cg.gpio_pin_expression(config[CONF_CS2_PIN])))

    spi_parent = await cg.get_variable(config[CONF_SPI_ID])
    cg.add(var.set_spi_parent(spi_parent))

    await display.register_display(var, config)