#!/bin/python3


from dataclasses import dataclass
import re
import sys


@dataclass
class StateMachine:
    name: str
    code: str
    states: list[str]
    events: list[str]
    transitions: list[tuple[str, str, str]]


def main():
    args = parse_args()
    for name, state_machine_txt in find_state_machines(args.input):
        stm = parse_state_machine(state_machine_txt, name)
        print(f"# {stm.name}")
        print(state_machine_txt)

        stm_simple_name = stm.name
        if stm_simple_name.endswith("_t"):
            stm_simple_name = stm_simple_name[:-2]
        headerfile = f"{args.header}/{stm_simple_name}_statemachine.hpp"
        cpp_code = generate_cpp_header(stm)
        write_cpp_code(cpp_code, headerfile)

        sourcefile = f"{args.source}/{stm_simple_name}_statemachine.cpp"
        cpp_code = generate_cpp_source(stm)
        write_cpp_code(cpp_code, sourcefile)

        print("Generated files:")
        print(f"    Header: {headerfile}")
        print(f"    Source: {sourcefile}")
        print()


def find_state_machines(input_file):
    with open(input_file, "r") as f:
        lines = f.readlines()
    state_machine_txt = ""
    name = ""
    in_state_machine = False
    for line in lines:
        if line.startswith("# "):
            if name and state_machine_txt:
                yield name, state_machine_txt
            state_machine_txt = ""
            name = line[1:].strip()
        if in_state_machine:
            if line.startswith("```"):
                in_state_machine = False
            else:
                state_machine_txt += line
        if line.startswith("```mermaid"):
            in_state_machine = True

    if name and state_machine_txt:
        yield name, state_machine_txt


RE_TRANSITION = re.compile(r"(\w+) --> (\w+)(: (\w+)(\((\w+)\))?)?")


def parse_state_machine(state_machine_txt, name):
    lines = state_machine_txt.split("\n")[1:]
    transitions = []
    states = ["Error"]
    events = list()
    for line in lines:
        line = line.strip()
        if not line:
            continue
        if line[0] == "#":
            continue

        match = RE_TRANSITION.match(line)
        if match:
            source, target, _, event, _, guard = match.groups()
            if event not in events:
                events.append(event)
            if source not in states:
                states.append(source)
            if target not in states:
                states.append(target)
            transitions.append((source, target, event))
        else:
            print(f"Invalid line: {line}")
            sys.exit(1)

    return StateMachine(
        states=states,
        events=events,
        transitions=transitions,
        name=name,
        code=state_machine_txt,
    )


def generate_cpp_header(stm):
    code = "// This is an autogenerated file. It should be included at header\n"
    code += "public:\n"
    code += "/// for the header / class\n\n"
    code += "enum state_e{\n"
    for state in stm.states:
        code += f"    {state},\n"
    code += "};\n\n"
    code += "static const char *to_string(state_e state);\n\n"

    code += "enum event_e{\n"
    for event in stm.events:
        code += f"    {event},\n"
    code += "};\n\n"
    code += "static const char *to_string(event_e event);\n\n"

    code += "state_e state = state_e::" + list(stm.states)[1] + ";\n\n"
    code += "void handle_event(event_e event);\n\n"

    code += "protected:\n"
    # A function for each state
    for state in stm.states:
        snake_name = re.sub(r"([a-z])([A-Z])", r"\1_\2", state).lower()
        code += f"void state_{snake_name}();\n"

    return code


def generate_cpp_source(stm):
    code = ""
    code += "// This is an autogenerated file. It should be included at source\n"
    code += "/** STATE MACHINE\n"
    code += " * ```mermaid\n"
    code += cpp_comment(stm.code)
    code += " * ```\n"
    code += " */\n"

    code += f"void {stm.name}::handle_event(event_e event){{\n"
    code += "    state_e oldstate = state;\n"
    code += "    state = state_e::Error;\n"
    code += "    switch(oldstate){\n"
    for state in stm.states:
        code += f"        case state_e::{state}:\n"
        code += "            switch(event){\n"
        for transition in stm.transitions:
            source, target, event = transition
            if source == state:
                code += f"                case event_e::{event}:\n"
                code += f"                    state = state_e::{target};\n"
                code += "                    break;\n"
        code += "            default:;\n"
        code += "            }\n"
        code += "            break;\n"
    code += "    }\n"

    # call a function with the state name in snake_case
    code += "    switch(state){\n"
    for state in stm.states:
        snake_name = re.sub(r"([a-z])([A-Z])", r"\1_\2", state).lower()
        code += f"        case state_e::{state}:\n"
        if state == "Error":
            code += '            DEBUG("Error, previous state: {}", oldstate);\n'
        code += f"            state_{snake_name}();\n"
        code += "            break;\n"
    code += "    }\n"

    code += "}\n"
    # formatters
    code += f"""// event formatter
const char *{stm.name}::to_string({stm.name}::event_e value) {{
    switch (value) {{
"""
    for event in stm.events:
        code += f"    case {stm.name}::event_e::{event}:\n"
        code += f'        return "{event}";\n'
    code += """    default:
            return "unknown";
"""
    code += "    }  \n"
    code += "}\n"

    # state formatter
    code += f"""// state formatter
const char *{stm.name}::to_string({stm.name}::state_e value) {{
    switch (value) {{
"""
    for state in stm.states:
        code += f"    case {stm.name}::state_e::{state}:\n"
        code += f'        return  "{state}";\n'
    code += """    default:"""
    code += '        return  "unknown";\n'
    code += "    }\n"
    code += "}\n"

    return code


def write_cpp_code(cpp_code, output_file):
    if output_file == "-":
        print(cpp_code)
    else:
        with open(output_file, "wt") as f:
            f.write(cpp_code)


def cpp_comment(text: str):
    """
    Returns the text as nice cpp comment lines
    """
    lines = text.split("\n")
    code = " * " + "\n * ".join(lines) + "\n"
    return code


def parse_args():
    import argparse

    parser = argparse.ArgumentParser(description="Convert a statemachine to C++ code")
    parser.add_argument("input", type=str, help="Path to the statemachine file .md")
    # output is optional, default is stdout
    parser.add_argument(
        "--header",
        nargs="?",
        type=str,
        help="Path to the output hpp file (- for stdout)",
        default="-",
    )
    parser.add_argument(
        "--source",
        nargs="?",
        type=str,
        help="Path to the output cpp file (- for stdout)",
        default="-",
    )
    return parser.parse_args()


if __name__ == "__main__":
    main()
