import os
import sys

from typing import List, Any, Collection, Mapping, Iterator, NamedTuple, cast
from typing import Dict, Tuple, Optional, Deque
from contextlib import contextmanager
from collections import deque

from google.protobuf.descriptor_pb2 import FileDescriptorProto, DescriptorProto
from google.protobuf.compiler.plugin_pb2 import CodeGeneratorRequest
from google.protobuf.compiler.plugin_pb2 import CodeGeneratorResponse

from .. import const
from .. import client
from .. import server


SUFFIX = '_grpc.py'

_CARDINALITY = {
    (False, False): const.Cardinality.UNARY_UNARY,
    (True, False): const.Cardinality.STREAM_UNARY,
    (False, True): const.Cardinality.UNARY_STREAM,
    (True, True): const.Cardinality.STREAM_STREAM,
}


class Method(NamedTuple):
    name: str
    cardinality: const.Cardinality
    request_type: str
    reply_type: str


class Service(NamedTuple):
    name: str
    methods: List[Method]


class Buffer:

    def __init__(self) -> None:
        self._lines: List[str] = []
        self._indent = 0

    def add(
        self,
        string: str,
        *args: Collection[Any], **kwargs: Mapping[str, Any],
    ) -> None:
        line = ' ' * self._indent * 4 + string.format(*args, **kwargs)
        self._lines.append(line.rstrip(' '))

    @contextmanager
    def indent(self) -> Iterator[None]:
        self._indent += 1
        try:
            yield
        finally:
            self._indent -= 1

    def content(self) -> str:
        return '\n'.join(self._lines) + '\n'


def render(
    proto_file: str,
    package: str,
    imports: Collection[str],
    services: Collection[Service],
) -> str:
    buf = Buffer()
    buf.add('# Generated by the Protocol Buffers compiler. DO NOT EDIT!')
    buf.add('# source: {}', proto_file)
    buf.add('# plugin: {}', __name__)
    if not services:
        return buf.content()

    buf.add('import abc')
    buf.add('import typing')
    buf.add('')
    buf.add('import {}', const.__name__)
    buf.add('import {}', client.__name__)
    buf.add('if typing.TYPE_CHECKING:')
    with buf.indent():
        buf.add('import {}', server.__name__)

    buf.add('')
    for mod in imports:
        buf.add('import {}', mod)
    for service in services:
        if package:
            service_name = '{}.{}'.format(package, service.name)
        else:
            service_name = service.name
        buf.add('')
        buf.add('')
        buf.add('class {}Base(abc.ABC):', service.name)
        with buf.indent():
            for (name, _, request_type, reply_type) in service.methods:
                buf.add('')
                buf.add('@abc.abstractmethod')
                buf.add("async def {}(self, stream: '{}.{}[{}, {}]') -> None:",
                        name, server.__name__, server.Stream.__name__,
                        request_type, reply_type)
                with buf.indent():
                    buf.add('pass')
            buf.add('')
            buf.add('def __mapping__(self) -> typing.Dict[str, {}.{}]:',
                    const.__name__,
                    const.Handler.__name__)
            with buf.indent():
                buf.add('return {{')
                with buf.indent():
                    for method in service.methods:
                        name, cardinality, request_type, reply_type = method
                        full_name = '/{}/{}'.format(service_name, name)
                        buf.add("'{}': {}.{}(", full_name, const.__name__,
                                const.Handler.__name__)
                        with buf.indent():
                            buf.add('self.{},', name)
                            buf.add('{}.{}.{},', const.__name__,
                                    const.Cardinality.__name__,
                                    cardinality.name)
                            buf.add('{},', request_type)
                            buf.add('{},', reply_type)
                        buf.add('),')
                buf.add('}}')

        buf.add('')
        buf.add('')
        buf.add('class {}Stub:', service.name)
        with buf.indent():
            buf.add('')
            buf.add('def __init__(self, channel: {}.{}) -> None:'
                    .format(client.__name__, client.Channel.__name__))
            with buf.indent():
                for method in service.methods:
                    name, cardinality, request_type, reply_type = method
                    full_name = '/{}/{}'.format(service_name, name)
                    method_cls: type
                    if cardinality is const.Cardinality.UNARY_UNARY:
                        method_cls = client.UnaryUnaryMethod
                    elif cardinality is const.Cardinality.UNARY_STREAM:
                        method_cls = client.UnaryStreamMethod
                    elif cardinality is const.Cardinality.STREAM_UNARY:
                        method_cls = client.StreamUnaryMethod
                    elif cardinality is const.Cardinality.STREAM_STREAM:
                        method_cls = client.StreamStreamMethod
                    else:
                        raise TypeError(cardinality)
                    method_cls = cast(type, method_cls)  # FIXME: redundant
                    buf.add('self.{} = {}.{}('.format(name, client.__name__,
                                                      method_cls.__name__))
                    with buf.indent():
                        buf.add('channel,')
                        buf.add('{!r},'.format(full_name))
                        buf.add('{},', request_type)
                        buf.add('{},', reply_type)
                    buf.add(')')
    return buf.content()


def _get_proto(request: CodeGeneratorRequest, name: str) -> FileDescriptorProto:
    return next(f for f in request.proto_file if f.name == name)


def _proto2py(proto_name: str) -> str:
    return proto_name.replace('/', '.')[:-len('.proto')] + '_pb2'


def _type_names(
    proto_file: FileDescriptorProto,
    message_type: DescriptorProto,
    parents: Optional[Deque[str]] = None,
) -> Iterator[Tuple[str, str]]:
    if parents is None:
        parents = deque()

    proto_name_parts = ['']
    if proto_file.package:
        proto_name_parts.append(proto_file.package)
    proto_name_parts.extend(parents)
    proto_name_parts.append(message_type.name)

    py_name_parts = [_proto2py(proto_file.name)]
    py_name_parts.extend(parents)
    py_name_parts.append(message_type.name)

    yield '.'.join(proto_name_parts), '.'.join(py_name_parts)

    parents.append(message_type.name)
    for nested in message_type.nested_type:
        yield from _type_names(proto_file, nested, parents=parents)
    parents.pop()


def main() -> None:
    with os.fdopen(sys.stdin.fileno(), 'rb') as inp:
        request = CodeGeneratorRequest.FromString(inp.read())

    types_map: Dict[str, str] = {}
    for pf in request.proto_file:
        for mt in pf.message_type:
            types_map.update(_type_names(pf, mt))

    response = CodeGeneratorResponse()
    for file_to_generate in request.file_to_generate:
        proto_file = _get_proto(request, file_to_generate)

        imports = [_proto2py(dep)
                   for dep in list(proto_file.dependency) + [file_to_generate]]

        services = []
        for service in proto_file.service:
            methods = []
            for method in service.method:
                cardinality = _CARDINALITY[(method.client_streaming,
                                            method.server_streaming)]
                methods.append(Method(
                    name=method.name,
                    cardinality=cardinality,
                    request_type=types_map[method.input_type],
                    reply_type=types_map[method.output_type],
                ))
            services.append(Service(name=service.name,
                                    methods=methods))

        file = response.file.add()
        file.name = file_to_generate.replace('.proto', SUFFIX)
        file.content = render(
            proto_file=proto_file.name,
            package=proto_file.package,
            imports=imports,
            services=services,
        )

    with os.fdopen(sys.stdout.fileno(), 'wb') as out:
        out.write(response.SerializeToString())
