from grpclib.plugin.main import Buffer

from ..config import ConfigSpec


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)
