import asyncio
import importlib.util
import inspect
import os
import signal
import sys
import time
import types
from concurrent import futures
from posixpath import basename
from typing import Callable, Type

import grpc
from grpc_health.v1 import health_pb2, health_pb2_grpc
from grpc_health.v1.health import HealthServicer
from grpc_reflection.v1alpha import reflection
from grpc_tools import protoc
from pydantic import BaseModel


def connect_obj_with_stub(pb2_grpc_module, pb2_module, service_obj: object) -> type:
    """
    Dynamically creates a concrete service class that inherits from a stub class generated by protoc, and adds stub methods
    to the concrete service class that call the corresponding methods on the original service object.

    Args:
        pb2_grpc_module: The generated gRPC module that contains the stub classes.
        pb2_module: The generated protobuf module that contains the request and response message classes.
        obj: The original service object that contains the methods to be called by the stub methods.

    Returns:
        A concrete service class instance that inherits from a stub class generated by protoc, and has stub methods that call the
        corresponding methods on the original service object.
    """
    # The stub class name is ServiceName + "Servicer" as per gRPC convention
    # For example, if ServiceName is Greeter, the stub class name will be GreeterServicer
    service_class = service_obj.__class__
    stub_class_name = service_class.__name__ + "Servicer"

    stub_class = getattr(pb2_grpc_module, stub_class_name)

    # Inherit the stub class and implement a concrete service class
    class ConcreteServiceClass(stub_class):
        pass

    def implement_stub_method(method):
        sig = inspect.signature(method)
        arg_type = get_request_arg_type(sig)
        converter = generate_message_converter(arg_type)
        response_type = sig.return_annotation
        size_of_parameters = len(sig.parameters)
        match size_of_parameters:
            case 1:

                def stub_method1(self, request, context, method=method):
                    arg = converter(request)
                    return getattr(pb2_module, response_type.__name__)(
                        **method(arg).model_dump()
                    )

                return stub_method1
            case 2:

                def stub_method2(self, request, context, method=method):
                    arg = converter(request)
                    return getattr(pb2_module, response_type.__name__)(
                        **method(arg, context).model_dump()
                    )

                return stub_method2
            case _:
                raise Exception("Method must have exactly one or two parameters")

    # Implement methods corresponding to each method of the stub class
    for method_name, method in get_rpc_methods(service_obj):
        a_method = implement_stub_method(method)
        setattr(ConcreteServiceClass, method_name, a_method)

    return ConcreteServiceClass


def connect_obj_with_stub_async(pb2_grpc_module, pb2_module, obj: object) -> type:
    """
    Dynamically creates a concrete service class that inherits from a stub class generated by protoc, and adds stub methods
    to the concrete service class that call the corresponding methods on the original service object.
    This is a version of connect_obj_with_stub for AsyncIOServer.

    Args:
        pb2_grpc_module: The generated gRPC module that contains the stub classes.
        pb2_module: The generated protobuf module that contains the request and response message classes.
        obj: The original service object that contains the methods to be called by the stub methods.

    Returns:
        A concrete service class instance that inherits from a stub class generated by protoc, and has stub methods that call the
        corresponding methods on the original service object.
    """
    # The stub class name is ServiceName + "Servicer" as per gRPC convention
    # For example, if ServiceName is Greeter, the stub class name will be GreeterServicer
    service_class = obj.__class__
    stub_class_name = service_class.__name__ + "Servicer"

    stub_class = getattr(pb2_grpc_module, stub_class_name)

    # Inherit the stub class and implement a concrete service class装
    class ConcreteServiceClass(stub_class):
        pass

    def implement_stub_method(method):
        sig = inspect.signature(method)
        arg_type = get_request_arg_type(sig)
        converter = generate_message_converter(arg_type)
        response_type = sig.return_annotation
        size_of_parameters = len(sig.parameters)
        match size_of_parameters:
            case 1:

                async def stub_method1(self, request, context, method=method):
                    arg = converter(request)
                    return getattr(pb2_module, response_type.__name__)(
                        **(await method(arg)).model_dump()
                    )

                return stub_method1
            case 2:

                async def stub_method2(self, request, context, method=method):
                    arg = converter(request)
                    return getattr(pb2_module, response_type.__name__)(
                        **(await method(arg, context)).model_dump()
                    )

                return stub_method2
            case _:
                raise Exception("Method must have exactly one or two parameters")

    # Implement methods corresponding to each method of the stub class
    for method_name, method in get_rpc_methods(obj):
        a_method = implement_stub_method(method)
        setattr(ConcreteServiceClass, method_name, a_method)

    return ConcreteServiceClass


class Message(BaseModel):
    def __doc__(self):
        return ""


def primitiveProtoValueToPythonValue(value):
    return value


def generate_converter(annotation: Type) -> Callable:
    """
    Generates a function to convert a protobuf value to a Python value based on the given annotation.

    Args:
        annotation: The type annotation to use for the conversion.

    Returns:
        A function that takes a protocol buffer value and returns the converted Python value.
    """
    if annotation in (int, str, bool, bytes, float):
        return primitiveProtoValueToPythonValue
    elif isinstance(annotation, types.GenericAlias):
        if annotation.__origin__ is list:
            item_converter = generate_converter(annotation.__args__[0])

            def list_converter(value):
                return [item_converter(v) for v in value]

            return list_converter
        elif annotation.__origin__ is dict:
            key_converter = generate_converter(annotation.__args__[0])
            value_converter = generate_converter(annotation.__args__[1])

            def dict_converter(value):
                return {key_converter(k): value_converter(v) for k, v in value.items()}

            return dict_converter
    elif issubclass(annotation, Message):
        return generate_message_converter(annotation)

    raise Exception(f"annotation {annotation} is not supported")


def generate_message_converter(arg_type: Type[Message]):
    if arg_type is None or not issubclass(arg_type, Message):
        raise TypeError("Request arg must be subclass of Message")

    fields = arg_type.model_fields
    converters = {
        field: generate_converter(field_type.annotation)
        for field, field_type in fields.items()
    }

    def converter(request):
        rdict = {}
        for field in fields.keys():
            rdict[field] = converters[field](getattr(request, field))
        return arg_type(**rdict)

    return converter


def generate_grpc_code(proto_file, python_out, grpc_python_out) -> tuple | None:
    command = f"-I. --python_out={python_out} --grpc_python_out={grpc_python_out} {proto_file}"
    exit_code = protoc.main(command.split())

    if exit_code != 0:
        return None

    base = os.path.splitext(proto_file)[0]
    generated_pb2_file = f"{base}_pb2.py"
    generated_pb2_grpc_file = f"{base}_pb2_grpc.py"

    if python_out not in sys.path:
        sys.path.append(python_out)

    # Dynamically import stub code
    # _pb2_grpcy.py
    spec = importlib.util.spec_from_file_location(
        generated_pb2_grpc_file, os.path.join(python_out, generated_pb2_grpc_file)
    )
    if spec is None:
        return None
    pb2_grpc_module = importlib.util.module_from_spec(spec)

    if spec.loader is None:
        return None
    spec.loader.exec_module(pb2_grpc_module)

    # _pb2.py
    spec = importlib.util.spec_from_file_location(
        generated_pb2_file, os.path.join(python_out, generated_pb2_file)
    )
    if spec is None:
        return None
    pb2_module = importlib.util.module_from_spec(spec)

    if spec.loader is None:
        return None
    spec.loader.exec_module(pb2_module)

    return pb2_grpc_module, pb2_module


def to_pascal_case(name: str) -> str:
    return "".join(part.capitalize() for part in name.split("_"))


def get_rpc_methods(obj: object) -> list[tuple[str, types.MethodType]]:
    return [
        (to_pascal_case(attr_name), getattr(obj, attr_name))
        for attr_name in dir(obj)
        if inspect.ismethod(getattr(obj, attr_name))
    ]


def protobuf_type_mapping(python_type: Type) -> str | type | None:
    mapping = {
        int: "int32",
        str: "string",
        bool: "bool",
        bytes: "bytes",
        float: "float",
    }

    if hasattr(python_type, "__origin__"):
        if python_type.__origin__ is list:
            inner_type = python_type.__args__[0]
            inner_proto_type = protobuf_type_mapping(inner_type)
            if inner_proto_type:
                return f"repeated {inner_proto_type}"
        elif python_type.__origin__ is dict:
            key_type = python_type.__args__[0]
            value_type = python_type.__args__[1]
            key_proto_type = protobuf_type_mapping(key_type)
            value_proto_type = protobuf_type_mapping(value_type)
            if key_proto_type and value_proto_type:
                return f"map<{key_proto_type}, {value_proto_type}>"
    elif issubclass(python_type, Message):
        return python_type
    else:
        return mapping.get(python_type, None)

    return None


def comment_out(docstr: str) -> tuple[str]:
    if docstr is None:
        return tuple()
    return tuple(f"//" if line == "" else f"// {line}" for line in docstr.split("\n"))


def indent_lines(lines, indentation="    "):
    return "\n".join(indentation + line for line in lines)


def generate_message_definition(message_type: Type) -> tuple[str, list[type]]:
    fields = []
    refs = []
    index = 0
    for field_name, field in message_type.__annotations__.items():
        proto_typename = protobuf_type_mapping(field)
        if proto_typename is None:
            raise Exception(f"Type {field} is not supported")
        if not isinstance(proto_typename, str):
            proto_typename = proto_typename.__name__
            refs.append(field)
        fields.append(f"{proto_typename} {field_name} = {index + 1};")
        index += 1

    return (
        f"""message {message_type.__name__} {{
{indent_lines(fields)}
}}""",
        refs,
    )


def is_subclass_of_message(cls: type) -> bool:
    return "fastserve.core.Message" in [
        t.__module__ + "." + t.__name__ for t in cls.__bases__
    ]


def generate_proto(obj: object, package_name: str = "") -> str:
    """
    Generates a Protocol Buffer definition file for a given service object.

    Args:
        obj (object): The service object to generate the Protocol Buffer definition for.
        package_name (str, optional): The name of the package to use in the Protocol Buffer definition. Defaults to "".

    Returns:
        str: The generated Protocol Buffer definition.
    """
    service_class = obj.__class__
    service_name = service_class.__name__
    service_docstr = inspect.getdoc(service_class)
    if service_docstr is None:
        service_comment = ""
    else:
        service_comment = "\n".join(comment_out(service_docstr))

    rpc_definitions = []
    message_definitions = []
    for method_name, method in get_rpc_methods(obj):
        method_sig = inspect.signature(method)
        method_docstr = inspect.getdoc(method)
        request_type = get_request_arg_type(method_sig)
        response_type = method_sig.return_annotation

        if not is_subclass_of_message(request_type) or not is_subclass_of_message(
            response_type
        ):
            raise TypeError("Request and Response must be subclasses of Message")

        request_message_name = request_type.__name__
        response_message_name = response_type.__name__

        message_types = [None, request_type, response_type]
        done_types = set()
        while message_types:
            message_type = message_types.pop()
            if message_type in done_types:
                continue
            done_types.add(message_type)

            if message_type is None:
                break
            message_docstr = inspect.getdoc(message_type)
            if message_docstr is not None:
                for comment_line in comment_out(message_docstr):
                    message_definitions.append(comment_line)
            message_def, refs = generate_message_definition(message_type)
            message_definitions.append(message_def)
            message_definitions.append("")
            message_types.extend(refs)

        if method_docstr is not None:
            for comment_line in comment_out(method_docstr):
                rpc_definitions.append(comment_line)
        rpc_definitions.append(
            f"rpc {method_name} ({request_message_name}) returns ({response_message_name});"
        )

    if package_name == "":
        if service_name.endswith("Service"):
            package_name = service_name[: -len("Service")]
        else:
            package_name = service_name
        package_name = package_name.lower() + ".v1"

    return f"""syntax = "proto3";

package {package_name};

{service_comment}
service {service_name} {{
{indent_lines(rpc_definitions)}
}}

{indent_lines(message_definitions, indentation="")}
"""


def get_request_arg_type(sig):
    num_of_params = len(sig.parameters)
    if not (num_of_params == 1 or num_of_params == 2):
        raise Exception("Method must have exactly one or two parameters")
    return tuple(sig.parameters.values())[0].annotation


def generate_and_compile_proto(obj: object, package_name: str = ""):
    klass = obj.__class__
    # proto_file = generate_proto(klass)
    proto_file = generate_proto(obj, package_name)

    # Generate proto file
    proto_file_name = klass.__name__.lower() + ".proto"
    with open(proto_file_name, "w") as f:
        f.write(proto_file)

    # Generate gRPC stub code from protobuf
    gen = generate_grpc_code(proto_file_name, ".", ".")
    if gen is None:
        raise Exception("Failed to generate grpc code")
    return gen  # pb2_grpc_module, pb2_module


class Server:
    """
    A gRPC server that can mount and run multiple gRPC services.

    Args:
        max_workers (int): The maximum number of worker threads to use for handling incoming requests.
        interceptors (list): A list of interceptors to apply to all incoming requests.
    """

    def __init__(self, max_workers: int = 8, *interceptors) -> None:
        self._server = grpc.server(
            futures.ThreadPoolExecutor(max_workers), interceptors=interceptors
        )
        self._service_names = []
        self._package_name = ""
        self._port = 50051

    def set_package_name(self, package_name: str):
        """
        Sets the package name that will be written in the generated protobuf file.

        Args:
            package_name (str): The package name to use.
        """
        self._package_name = package_name

    def set_port(self, port: int):
        """
        Sets the port number to listen on for incoming requests.

        Args:
            port (int): The port number to use.
        """
        self._port = port

    def mount(self, obj: object, package_name: str = ""):
        """
        Mounts a service object as a gRPC service on the server.

        Args:
            obj (object): The gRPC service object to mount.
            package_name (str): The package name to use for generating gRPC service stubs.
        """
        pb2_grpc_module, pb2_module = generate_and_compile_proto(obj, package_name)
        self.mount_using_pb2_modules(pb2_grpc_module, pb2_module, obj)

    def mount_using_pb2_modules(self, pb2_grpc_module, pb2_module, obj: object):
        """
        Mounts a service object as a gRPC service on the server using pre-generated pb2_grpc and pb2 modules.

        Args:
            pb2_grpc_module: The pb2_grpc module to use for generating the gRPC service stub.
            pb2_module: The pb2 module to use for generating the gRPC service stub.
            obj (object): The gRPC service object to mount.
        """
        concreteServiceClass = connect_obj_with_stub(pb2_grpc_module, pb2_module, obj)

        service_name = obj.__class__.__name__
        service_impl = concreteServiceClass()

        # Add the servicer to gRPC server
        getattr(pb2_grpc_module, f"add_{service_name}Servicer_to_server")(
            service_impl, self._server
        )

        full_service_name = pb2_module.DESCRIPTOR.services_by_name[
            service_name
        ].full_name
        self._service_names.append(full_service_name)

    def run(self, *objs):
        """
        Runs the gRPC server with all mounted services.

        Args:
            *objs: The gRPC service objects to mount and run.
        """
        for obj in objs:
            self.mount(obj, self._package_name)

        # Register Helth Servcie and Reflection Service as well
        SERVICE_NAMES = (
            health_pb2.DESCRIPTOR.services_by_name["Health"].full_name,
            reflection.SERVICE_NAME,
            *self._service_names,
        )
        health_servicer = HealthServicer()
        health_pb2_grpc.add_HealthServicer_to_server(health_servicer, self._server)
        reflection.enable_server_reflection(SERVICE_NAMES, self._server)

        # Start gRPC server
        self._server.add_insecure_port(f"[::]:{self._port}")
        self._server.start()

        # Set up a signal handler to do graceful shutdown on receipt of SIGINT (Ctrl-C) or SIGTERM
        def handle_signal(signal, frame):
            print("Received shutdown signal...")
            self._server.stop(grace=10)
            print("gRPC server shutdown.")
            sys.exit(0)

        signal.signal(signal.SIGINT, handle_signal)
        signal.signal(signal.SIGTERM, handle_signal)

        print("gRPC server is running...")
        while True:
            time.sleep(86400)  # 1 day in seconds


class AsyncIOServer:
    """
    A gRPC server that can mount and run multiple gRPC services using asyncio.

    Args:
        interceptors (list): A list of interceptors to apply to all incoming requests.
    """

    def __init__(self, *interceptors) -> None:
        self._server = grpc.aio.server(interceptors=interceptors)
        self._service_names = []
        self._package_name = ""
        self._port = 50051

    def set_package_name(self, package_name: str):
        """
        Sets the package name that will be written in the generated protobuf file.

        Args:
            package_name (str): The package name to use.
        """
        self._package_name = package_name

    def set_port(self, port: int):
        """
        Sets the port number to listen on for incoming requests.

        Args:
            port (int): The port number to use.
        """
        self._port = port

    def mount(self, obj: object, package_name: str = ""):
        """
        Mounts a service object as a gRPC service on the server.

        Args:
            obj (object): The gRPC service object to mount.
            package_name (str): The package name to use for generating gRPC service stubs.
        """
        pb2_grpc_module, pb2_module = generate_and_compile_proto(obj, package_name)
        self.mount_using_pb2_modules(pb2_grpc_module, pb2_module, obj)

    def mount_using_pb2_modules(self, pb2_grpc_module, pb2_module, obj: object):
        """
        Mounts a service object as a gRPC service on the server using pre-generated pb2_grpc and pb2 modules.

        Args:
            pb2_grpc_module: The pb2_grpc module to use for generating the gRPC service stub.
            pb2_module: The pb2 module to use for generating the gRPC service stub.
            obj (object): The gRPC service object to mount.
        """
        concreteServiceClass = connect_obj_with_stub_async(
            pb2_grpc_module, pb2_module, obj
        )

        service_name = obj.__class__.__name__
        service_impl = concreteServiceClass()

        # Add the servicer to gRPC server
        getattr(pb2_grpc_module, f"add_{service_name}Servicer_to_server")(
            service_impl, self._server
        )

        full_service_name = pb2_module.DESCRIPTOR.services_by_name[
            service_name
        ].full_name
        self._service_names.append(full_service_name)

    async def run(self, *objs):
        """
        Runs the gRPC server with all mounted services.

        Args:
            *objs: The gRPC service objects to mount and run.
        """
        for obj in objs:
            self.mount(obj, self._package_name)

        # Register Helth Servcie and Reflection Service as well
        SERVICE_NAMES = (
            health_pb2.DESCRIPTOR.services_by_name["Health"].full_name,
            reflection.SERVICE_NAME,
            *self._service_names,
        )
        health_servicer = HealthServicer()
        health_pb2_grpc.add_HealthServicer_to_server(health_servicer, self._server)
        reflection.enable_server_reflection(SERVICE_NAMES, self._server)

        # Start gRPC server
        self._server.add_insecure_port(f"[::]:{self._port}")
        await self._server.start()

        shutdown_event = asyncio.Event()

        def shutdown(signum, frame):
            print("Received shutdown signal...")
            shutdown_event.set()

        # Set up a signal handler to do graceful shutdown on receipt of SIGINT (Ctrl-C) or SIGTERM
        for s in [signal.SIGTERM, signal.SIGINT]:
            signal.signal(s, shutdown)

        print("gRPC server is running...")
        await shutdown_event.wait()  # Wait for shutdown event to be set
        await self._server.stop(10)  # Wait for server to shutdown gracefully
        print("gRPC server shutdown.")


if __name__ == "__main__":
    py_file_name = sys.argv[1]
    class_name = sys.argv[2]
    module_name = os.path.splitext(basename(py_file_name))[0]
    module = importlib.import_module(module_name)
    klass = getattr(module, class_name)
    generate_and_compile_proto(klass())
