from betterosi.io import Writer, read, MESSAGES_TYPE  # noqa: F401
from google.protobuf.descriptor_pb2 import (
    EnumDescriptorProto,
    DescriptorProto,
    FileDescriptorProto,
)
from . import generated
from .generated.osi3 import *  # noqa: F403
from enum import EnumMeta
from google.protobuf.descriptor_pool import DescriptorPool


for c_name in generated.osi3.__all__:
    c = getattr(generated.osi3, c_name)
    if hasattr(c, "parse"):
        c.ParseFromString = c.parse


def get_proto_dependencies(obj):
    if hasattr(obj, "_betterproto"):
        return set(
            [
                o
                for os in [
                    [v] + list(get_proto_dependencies(v))
                    for v in obj._betterproto.cls_by_field.values()
                    if hasattr(v, "_serialized_pb")
                ]
                for o in os
            ]
        )
    else:
        return []


def insert_into_hierarchy_dict(d, name, cls):
    first = name.split(".")[0]
    if first not in d:
        d[first] = {"cls": None, "children": {}}
    if len(name.split(".")) > 1:
        insert_into_hierarchy_dict(
            d[first]["children"], ".".join(name.split(".")[1:]), cls
        )
    else:
        d[first]["cls"] = cls


def cls_to_proto(message, children):
    if isinstance(message, EnumMeta):
        proto = EnumDescriptorProto()
        proto.ParseFromString(message._serialized_pb())
        proto.name = proto.name.split(".")[-1]
        return None, proto
    else:
        proto = DescriptorProto()
        proto.ParseFromString(message._serialized_pb())
        proto.name = proto.name.split(".")[-1]
        proto.ClearField("enum_type")  # Remove all enums first
        proto.ClearField("nested_type")  # Remove all enums first
        for k, v in children.items():
            m, e = cls_to_proto(v["cls"], v["children"])
            if m is not None:
                proto.nested_type.add().MergeFrom(m)
            if e is not None:
                proto.enum_type.add().MergeFrom(e)
        return proto, None


def get_descriptor(cls, package="osi3"):
    fd = FileDescriptorProto(name=cls.__name__, package=package)

    package_hierarchy = {}
    for o in list(get_proto_dependencies(cls)) + [cls]:
        proto = DescriptorProto()
        proto.ParseFromString(o._serialized_pb())
        insert_into_hierarchy_dict(package_hierarchy, proto.name, o)

    messages = []
    enums = []
    for k, v in package_hierarchy.items():
        m, e = cls_to_proto(v["cls"], v["children"])
        if m is not None:
            messages.append(m)
        if e is not None:
            enums.append(e)
    fd = FileDescriptorProto(name=cls.__name__, package=package)
    fd.message_type.extend(messages)
    fd.enum_type.extend(enums)
    pool = DescriptorPool()
    pool.Add(fd)

    return pool.FindMessageTypeByName(f"{package}.{cls.__name__}")


for c_name in MESSAGES_TYPE:
    o = getattr(generated.osi3, c_name)
    setattr(o, "DESCRIPTOR", get_descriptor(o))
