import yaml

# CONFIGURATION
NR_OF_PAGES = 3
INCLUDE_DISPLAY_EXTRA = True  # Enable Double, Triple, Hold for main display buttons
LED_CONFIRMATION = True       # Flash LED on press
FILENAME = "3bar_1display_events.yaml"

# --- ESPHome YAML Handling Fixes ---
class ESPHomeExtend:
    def __init__(self, value): self.value = value

class ESPHomeLambda:
    def __init__(self, value): self.value = value

class LiteralString(str): pass

def represent_extend(dumper, data):
    return dumper.represent_scalar('!extend', data.value)

def represent_lambda(dumper, data):
    # Ensures !lambda is followed by a quoted string or block for ESPHome compatibility
    return dumper.represent_scalar('!lambda', data.value, style='"')

def represent_literal(dumper, data):
    return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')

yaml.add_representer(ESPHomeExtend, represent_extend)
yaml.add_representer(ESPHomeLambda, represent_lambda)
yaml.add_representer(LiteralString, represent_literal)
# -----------------------------------

def generate():
    # 1. Define event lists for substitutions
    base_types = ["click", "double_click", "triple_click", "hold"]
    display_types = [] + (base_types[1:] if INCLUDE_DISPLAY_EXTRA else [])
    
    display_events = []
    button_events = []
    for p in range(NR_OF_PAGES):
        display_events.extend([f"{t} page {p}" for t in display_types])
        button_events.extend([f"{t} page {p}" for t in base_types])

    data = {
        "substitutions": {
            "nr_of_pages": str(NR_OF_PAGES),
            "display_events": display_events,
            "button_events": button_events
        },
        "script": [{
            "id": "trigger_with_feedback",
            "parameters": {
                "event_id": "esphome::template_::TemplateEvent*",
                "event_msg": "std::string",
                "led_id": "light::LightState*",
                "r": "float", "g": "float", "b": "float"
            },
            "then": [
                {"lambda": "event_id->trigger(event_msg);"},
                {"if": {
                    "condition": {"lambda": "return led_id != nullptr;"},
                    "then": [{"lambda": LiteralString(
                        "auto call = led_id->make_call();\n"
                        "call.set_state(true);\n"
                        "call.set_rgb(r, g, b);\n"
                        "call.set_brightness(1.0);\n"
                        "call.set_flash_length(200); // Returns to previous state automatically\n"
                        "call.perform();"
                    )}]
                }}
            ]
        }],
        "event": []
    }

    # 2. Event Entities
    buttons = [
        ("DISPLAY Button Left", "display_btn_left_event", "${display_events}"),
        ("DISPLAY Button Right", "display_btn_right_event", "${display_events}"),
        ("BAR1 Button Left", "bar1_btn_left_event", "${button_events}"),
        ("BAR1 Button Right", "bar1_btn_right_event", "${button_events}"),
        ("BAR2 Button Left", "bar2_btn_left_event", "${button_events}"),
        ("BAR2 Button Right", "bar2_btn_right_event", "${button_events}"),
        ("BAR3 Button Left", "bar3_btn_left_event", "${button_events}"),
        ("BAR3 Button Right", "bar3_btn_right_event", "${button_events}"),
    ]
    for name, eid, types in buttons:
        data["event"].append({
            "platform": "template", "name": name, "id": eid,
            "device_class": "button", "event_types": types
        })

    # 3. Binary Sensor Extensions
    data["binary_sensor"] = []
    triggers = {
        "click": (0.0, 1.0, 0.0),        # Green
        "double_click": (0.0, 0.0, 1.0), # Blue
        "triple_click": (1.0, 0.0, 1.0), # Purple
        "hold": (1.0, 0.5, 0.0)          # Orange
    }

    btn_configs = [
        ("main_display_btn_left", "display_btn_left_event", None, "prev_page"),
        ("main_display_btn_right", "display_btn_right_event", None, "next_page"),
        ("bar1_btn_left", "bar1_btn_left_event", "bar1_left_rgb_front", None),
        ("bar1_btn_right", "bar1_btn_right_event", "bar1_right_rgb_front", None),
        ("bar2_btn_left", "bar2_btn_left_event", "bar2_left_rgb_front", None),
        ("bar2_btn_right", "bar2_btn_right_event", "bar2_right_rgb_front", None),
        ("bar3_btn_left", "bar3_btn_left_event", "bar3_left_rgb_front", None),
        ("bar3_btn_right", "bar3_btn_right_event", "bar3_right_rgb_front", None),
    ]

    for btn_id, event_id, led_id, page_script in btn_configs:
        sensor = {
            "id": ESPHomeExtend(btn_id),
            "filters": [{"delayed_off": "10ms"}],
            "internal": True,
            "on_multi_click": []
        }
        
        is_display = "main_display" in btn_id
        for t_name, color in triggers.items():
            if is_display and t_name != "click" and not INCLUDE_DISPLAY_EXTRA:
                continue
                
            multi_click = {"timing": []}
            if t_name == "click":
                multi_click["timing"] = ["ON for at most 300ms", "OFF for at least 0.2s"]
            elif t_name == "double_click":
                multi_click["timing"] = ["ON for at most 300ms", "OFF for at most 0.2s", "ON for at most 300ms", "OFF for at least 0.2s"]
            elif t_name == "triple_click":
                multi_click["timing"] = ["ON for at most 300ms", "OFF for at most 0.2s", "ON for at most 300ms", "OFF for at most 0.2s", "ON for at most 300ms", "OFF for at least 0.2s"]
            elif t_name == "hold":
                multi_click["timing"] = ["ON for at least 1s", "OFF for at least 0.2s"]

            actions = []
            if t_name == "click" and page_script:
                actions.append({"script.execute": page_script})
            else:
                # The Feedback call with proper ESPHome formatting
                actions.append({
                    "script.execute": {
                        "id": "trigger_with_feedback",
                        "event_id": event_id,
                        "event_msg": ESPHomeLambda(f'return "{t_name} page " + std::to_string(id(actual_page));'),
                        "led_id": led_id if (led_id and LED_CONFIRMATION) else ESPHomeLambda("return nullptr;"),
                        "r": color[0], "g": color[1], "b": color[2]
                    }
                })
            multi_click["then"] = actions
            sensor["on_multi_click"].append(multi_click)
        data["binary_sensor"].append(sensor)

    # 4. Write output
    with open(FILENAME, "w") as f:
        f.write("## Dynamically generated event configuration\n")
        f.write("## Do not edit manually - run generator script to update\n\n")
        yaml.dump(data, f, sort_keys=False, default_flow_style=False)
    print(f"Generated {FILENAME} successfully.")

if __name__ == "__main__":
    generate()