from ..config import ConfigSpec
from ..runtime._utils import Buffer


def _render_wires(proto_file: str, spec: ConfigSpec):
    buf = Buffer()
    buf.add("# Generated by the Protocol Buffers compiler. DO NOT EDIT!")
    buf.add(f"# source: {proto_file}")
    buf.add(f"# plugin: {__name__}")

    if any(wire.optional for wire in spec.wires):
        buf.add("from typing import Optional")

    buf.add("from dataclasses import dataclass")
    buf.add("")

    imports = set()
    for wire in spec.wires:
        wire_type = wire.value.WhichOneof("type")
        if wire_type == "input":
            value = wire.value.input.type
        elif wire_type == "output":
            value = wire.value.output.type
        else:
            raise NotImplementedError(wire_type)
        module_name, _ = value.rsplit(".", 1)
        imports.add(module_name)
    for module_name in sorted(imports):
        buf.add(f"import {module_name}")
    buf.add("")
    buf.add("")

    buf.add("@dataclass")
    buf.add(f"class WiresIn:")
    with buf.indent():
        empty = True
        for wire in spec.wires:
            if wire.value.WhichOneof("type") == "input":
                if wire.optional:
                    wire_type = f"Optional[{wire.value.input.type}]"
                else:
                    wire_type = wire.value.input.type
                buf.add(f"{wire.name}: {wire_type}")
                empty = False
        if empty:
            buf.add("pass")
    buf.add("")
    buf.add("")
    buf.add("@dataclass")
    buf.add(f"class WiresOut:")
    with buf.indent():
        empty = True
        for wire in spec.wires:
            if wire.value.WhichOneof("type") == "output":
                buf.add(f"{wire.name}: {wire.value.output.type}")
                empty = False
        if empty:
            buf.add("pass")
    file_name = proto_file.replace(".proto", "_wires.py")
    return file_name, buf.content()


def _render_entrypoint(proto_file: str):
    buf = Buffer()
    buf.add("# Generated by the Protocol Buffers compiler. DO NOT EDIT!")
    buf.add(f"# source: {proto_file}")
    buf.add(f"# plugin: {__name__}")
    buf.add("import sys")
    buf.add("")
    buf.add("from harness.runtime import Runner")
    buf.add("")

    main_module = proto_file.replace("/", ".").replace(".proto", "")
    pb2_module = proto_file.replace("/", ".").replace(".proto", "_pb2")
    wires_module = proto_file.replace("/", ".").replace(".proto", "_wires")

    buf.add(f"import {main_module}")
    buf.add(f"import {pb2_module}")
    buf.add(f"import {wires_module}")
    buf.add("")
    buf.add("runner = Runner(")
    buf.add(f"    {pb2_module}.Configuration,")
    buf.add(f"    {wires_module}.WiresIn,")
    buf.add(f"    {wires_module}.WiresOut,")
    buf.add(")")
    buf.add("")
    buf.add("if __name__ == '__main__':")
    with buf.indent():
        buf.add(f"sys.exit(runner.run({main_module}.setup, sys.argv))")
    return "entrypoint.py", buf.content()


def render(proto_file: str, spec: ConfigSpec):
    yield _render_wires(proto_file, spec)
    yield _render_entrypoint(proto_file)
