import struct
import socket
from .packing_tools import bytes_to_short, bytes_to_integer
from .ip import IP4Prefix, IP4Address
from .ip import IP6Prefix, IP6Address
from io import BytesIO


class BgpMessage(object):
    OPEN_MESSAGE = 1
    UPDATE_MESSAGE = 2
    NOTIFICATION_MESSAGE = 3
    KEEPALIVE_MESSAGE = 4
    MARKER = b"\xFF" * 16
    HEADER_LENGTH = 19


PARSERS = {}


class BgpMessageParser(object):
    def __init__(self):
        self.capabilities = {}

    def parse(self, message_type, serialised_message):
        return PARSERS[message_type](serialised_message, self.capabilities)


def register_parser(cls):
    PARSERS[cls.MSG_TYPE] = cls.parse
    return cls


class BgpMessagePacker(object):
    def __init__(self):
        self.capabilities = {}

    def pack(self, message):
        packed_message = message.pack(self.capabilities)
        length = BgpMessage.HEADER_LENGTH + len(packed_message)
        header = struct.pack("!16sHB", BgpMessage.MARKER, length, message.MSG_TYPE)
        return header + packed_message


MULTIPROTOCOL_TYPES = {(1, 1): "ipv4-unicast", (2, 1): "ipv6-unicast"}


def parse_multiprotocol(serialised_capability):
    afi, reserved, safi = struct.unpack("!HBB", serialised_capability)
    if (afi, safi) in MULTIPROTOCOL_TYPES:
        return MULTIPROTOCOL_TYPES[(afi, safi)]
    return (afi, safi)


def parse_routerefresh(serialised_capability):
    return True


def parse_fourbyteas(serialised_capability):
    (peer_as,) = struct.unpack("!I", serialised_capability)
    return peer_as


capability_parsers = {
    1: parse_multiprotocol,
    2: parse_routerefresh,
    65: parse_fourbyteas,
}

capability_keys = {
    1: "multiprotocol",
    2: "routerefresh",
    65: "fourbyteas",
}


def parse_capabilities(serialised_capabilities):
    stream = BytesIO(serialised_capabilities)
    capabilities = {}

    while True:
        serialised_header = stream.read(2)
        if len(serialised_header) == 0:
            break
        capability_code, capability_length = struct.unpack("!BB", serialised_header)
        serialised_capability = stream.read(capability_length)
        if capability_code in capability_keys:
            capability_key = capability_keys[capability_code]
            if capability_key not in capabilities:
                capabilities[capability_key] = []
            capabilities[capability_key].append(
                capability_parsers[capability_code](serialised_capability)
            )
        else:
            print("WARNING did not recognise capability code %d" % capability_code)

    return capabilities


MULTIPROTOCOL_AFI_SAFI = {
    "ipv4-unicast": (1, 1),
    "ipv6-unicast": (2, 1),
}


def pack_multiprotocol(multiprotocol):
    afi, safi = MULTIPROTOCOL_AFI_SAFI[multiprotocol]
    return struct.pack("!HBB", afi, 0, safi)


def pack_routerefresh(route_refresh):
    return b""


def pack_fourbyteas(fourbyteas):
    return struct.pack("!I", fourbyteas)


capability_packers = {
    "multiprotocol": pack_multiprotocol,
    "routerefresh": pack_routerefresh,
    "fourbyteas": pack_fourbyteas,
}

capability_numbers = {
    "multiprotocol": 1,
    "routerefresh": 2,
    "fourbyteas": 65,
}

OPTIONAL_PARAMETER_CAPABILITY = 2


def pack_capabilities(capabilities):
    packed_capability_list = []

    for capability_key, capability_list in capabilities.items():
        capability_code = capability_numbers[capability_key]
        for capability in capability_list:
            packed_body = capability_packers[capability_key](capability)
            packed_header = struct.pack("!BB", capability_code, len(packed_body))
            packed_capability_list.append(packed_header + packed_body)

    packed_capabilities = b"".join(sorted(packed_capability_list))
    return packed_capabilities


def merge_dict_of_lists(main_dict, new_dict):
    for key, new_list in new_dict.items():
        if key not in main_dict:
            main_dict[key] = new_list
        else:
            main_dict[key] += new_list


def parse_optional_parameters(serialised_optional_parameters):
    stream = BytesIO(serialised_optional_parameters)
    capabilities = {}

    while True:
        serialised_header = stream.read(2)
        if len(serialised_header) == 0:
            break
        parameter_type, parameter_length = struct.unpack("!BB", serialised_header)
        if parameter_type != OPTIONAL_PARAMETER_CAPABILITY:
            raise ValueError(
                "OPEN: Got unsupported optional parameter: %d" % parameter_type
            )
        serialised_capabilities = stream.read(parameter_length)

        next_capabilities = parse_capabilities(serialised_capabilities)
        merge_dict_of_lists(capabilities, next_capabilities)

    return capabilities


@register_parser
class BgpOpenMessage(BgpMessage):
    MSG_TYPE = BgpMessage.OPEN_MESSAGE

    def __init__(self, version, peer_as, hold_time, identifier, capabilities):
        self.version = version
        self.peer_as = peer_as
        self.hold_time = hold_time
        self.identifier = identifier
        self.capabilities = capabilities

    @classmethod
    def parse(cls, serialised_message, _capabilities):
        (
            version,
            peer_as,
            hold_time,
            identifier,
            optional_parameters_length,
        ) = struct.unpack("!BHH4sB", serialised_message[:10])
        capabilities = parse_optional_parameters(
            serialised_message[10 : 10 + optional_parameters_length]
        )
        return cls(version, peer_as, hold_time, IP4Address(identifier), capabilities)

    def pack(self, _capabilities):
        packed_capabilities = pack_capabilities(self.capabilities)
        capabilities_header = struct.pack(
            "!BB", OPTIONAL_PARAMETER_CAPABILITY, len(packed_capabilities)
        )

        return (
            struct.pack(
                "!BHH4sB",
                self.version,
                self.peer_as,
                self.hold_time,
                self.identifier.address,
                len(capabilities_header + packed_capabilities),
            )
            + capabilities_header
            + packed_capabilities
        )

    def __str__(self):
        return (
            "BgpOpenMessage: Version %s, Peer AS: %s, Hold time: %s, Identifier: %s"
            % (self.version, self.peer_as, self.hold_time, self.identifier)
        )


IP4_LENGTH = 4


def prefix_byte_length(bit_length):
    byte_length = bit_length // 8
    if bit_length % 8:
        byte_length += 1

    return byte_length


def pack_prefix(prefix, length):
    return prefix[: prefix_byte_length(length)]


def unpack_prefix(prefix):
    num_extra_bytes = IP4_LENGTH - len(prefix)

    if num_extra_bytes == 0:
        return prefix

    return prefix + b"\x00" * num_extra_bytes


def parse_nlri(serialised_nlri):
    stream = BytesIO(serialised_nlri)
    prefixes = []

    while True:
        serialised_length = stream.read(1)
        if len(serialised_length) == 0:
            break
        prefix_length = ord(serialised_length)
        packed_prefix = stream.read(prefix_byte_length(prefix_length))
        prefix = unpack_prefix(packed_prefix)
        prefixes.append(IP4Prefix(prefix, prefix_length))

    return prefixes


ORIGIN_CODES = {0: "IGP", 1: "EGP", 2: "INCOMPLETE"}


def parse_origin(packed_origin):
    return ORIGIN_CODES[ord(packed_origin)]


AS_SET_CODE = 1
AS_SEQUENCE_CODE = 2
AS_NUMBER_LENGTH = 2
AS4_NUMBER_LENGTH = 4


def parse_as4_path(packed_as_path):
    # this does as_sets wrong, assumes everything is as_sequence
    input_stream = BytesIO(packed_as_path)
    as_numbers = []
    while True:
        packed_type_and_count = input_stream.read(2)
        if len(packed_type_and_count) == 0:
            break
        type_code, count = struct.unpack("!BB", packed_type_and_count)
        if type_code == AS_SET_CODE:
            print("WARNING received update with AS_SET, treating like AS_SEQUENCE")
        packed_as_sequence = input_stream.read(count * AS4_NUMBER_LENGTH)
        as_numbers += struct.unpack("!" + ("I" * count), packed_as_sequence)
    return " ".join(["%d" % x for x in as_numbers])


def parse_as_path(packed_as_path):
    # this does as_sets wrong, assumes everything is as_sequence
    input_stream = BytesIO(packed_as_path)
    as_numbers = []
    while True:
        packed_type_and_count = input_stream.read(2)
        if len(packed_type_and_count) == 0:
            break
        type_code, count = struct.unpack("!BB", packed_type_and_count)
        if type_code == AS_SET_CODE:
            print("WARNING received update with AS_SET, treating like AS_SEQUENCE")
        packed_as_sequence = input_stream.read(count * AS_NUMBER_LENGTH)
        as_numbers += struct.unpack("!" + ("H" * count), packed_as_sequence)
    return " ".join(["%d" % x for x in as_numbers])


def parse_next_hop(packed_next_hop):
    return IP4Address(packed_next_hop)


IP6_AFI = 2
UNICAST_SAFI = 1

IP6_LENGTH = 16


def unpack_prefix6(prefix):
    num_extra_bytes = IP6_LENGTH - len(prefix)

    if num_extra_bytes == 0:
        return prefix

    return prefix + b"\x00" * num_extra_bytes


def parse_nlri6(stream):
    prefixes = []

    while True:
        serialised_length = stream.read(1)
        if len(serialised_length) == 0:
            break
        prefix_length = ord(serialised_length)
        packed_prefix = stream.read(prefix_byte_length(prefix_length))
        prefix = unpack_prefix6(packed_prefix)
        prefixes.append(IP6Prefix(prefix, prefix_length))

    return prefixes


def parse_mp_reach_nlri(packed_mp_reach_nlri):
    attributes = {}
    stream = BytesIO(packed_mp_reach_nlri)
    afi, safi, next_hop_length = struct.unpack("!HBB", stream.read(4))
    if afi != IP6_AFI:
        raise ValueError("MP_REACH_NLRI: Got unsupported AFI: %d" % afi)
    if safi != UNICAST_SAFI:
        raise ValueError("MP_REACH_NLRI: Got unsupported SAFI: %d" % safi)
    if next_hop_length % IP6_LENGTH != 0:
        raise ValueError(
            "MP_REACH_NLRI: Got unsupported next hop length: %d" % next_hop_length
        )

    attributes["next_hop"] = []
    for _ in range(next_hop_length // 16):
        attributes["next_hop"].append(
            IP6Address(struct.unpack("!16s", stream.read(IP6_LENGTH))[0])
        )

    _reserved_snpa = stream.read(1)

    attributes["nlri"] = parse_nlri6(stream)

    return attributes


def parse_mp_unreach_nlri(packed_mp_reach_nlri):
    attributes = {}
    stream = BytesIO(packed_mp_reach_nlri)
    afi, safi = struct.unpack("!HB", stream.read(3))
    if afi != IP6_AFI:
        raise ValueError("MP_UNREACH_NLRI: Got unsupported AFI: %d" % afi)
    if safi != UNICAST_SAFI:
        raise ValueError("MP_UNREACH_NLRI: Got unsupported SAFI: %d" % safi)

    attributes["withdrawn_routes"] = parse_nlri6(stream)

    return attributes


attribute_parsers = {
    1: parse_origin,
    2: parse_as_path,
    3: parse_next_hop,
    14: parse_mp_reach_nlri,
    15: parse_mp_unreach_nlri,
    17: parse_as4_path,
}

attribute_keys = {
    1: "origin",
    2: "as_path",
    3: "next_hop",
    14: "mp_reach_nlri",
    15: "mp_unreach_nlri",
    17: "as4_path",
}

ORIGIN_NUMBERS = {"IGP": 0, "EGP": 1, "INCOMPLETE": 2}


def pack_origin(origin):
    return struct.pack("!B", ORIGIN_NUMBERS[origin])


def pack_as4_path(as_path):
    if not as_path:
        return b""

    as_numbers = [int(x) for x in as_path.split(" ")]
    count = len(as_numbers)
    header = struct.pack("!BB", AS_SEQUENCE_CODE, count)
    body = struct.pack("!" + ("I" * count), *as_numbers)
    return header + body


def pack_as_path(as_path):
    if not as_path:
        return b""

    as_numbers = [int(x) for x in as_path.split(" ")]
    count = len(as_numbers)
    header = struct.pack("!BB", AS_SEQUENCE_CODE, count)
    body = struct.pack("!" + ("H" * count), *as_numbers)
    return header + body


def pack_next_hop(next_hop):
    return next_hop.address


def pack_nlri6(nlri):
    packed_nlri = []

    for prefix in nlri:
        packed_prefix = struct.pack("!B", prefix.length) + pack_prefix(
            prefix.prefix, prefix.length
        )
        packed_nlri.append(packed_prefix)

    return b"".join(packed_nlri)


def pack_mp_reach_nlri(mp_reach_nlri):
    packed_header = struct.pack(
        "!HBB", IP6_AFI, UNICAST_SAFI, IP6_LENGTH * len(mp_reach_nlri["next_hop"])
    )
    packed_next_hop_list = []
    for next_hop in mp_reach_nlri["next_hop"]:
        packed_next_hop_list.append(struct.pack("!16s", next_hop.address))
    packed_next_hops = b"".join(packed_next_hop_list)
    packed_reserved = struct.pack("!B", 0)
    packed_nlri = pack_nlri6(mp_reach_nlri["nlri"])

    return packed_header + packed_next_hops + packed_reserved + packed_nlri


def pack_mp_unreach_nlri(mp_unreach_nlri):
    packed_header = struct.pack("!HB", IP6_AFI, UNICAST_SAFI)
    packed_withdrawn_route_list = []
    for withdrawn_route in mp_unreach_nlri["withdrawn_routes"]:
        packed_prefix = struct.pack("!B", withdrawn_route.length) + pack_prefix(
            withdrawn_route.prefix, withdrawn_route.length
        )
        packed_withdrawn_route_list.append(packed_prefix)
    packed_withdrawn_routes = b"".join(packed_withdrawn_route_list)

    return packed_header + packed_withdrawn_routes


attribute_packers = {
    "origin": pack_origin,
    "as_path": pack_as_path,
    "next_hop": pack_next_hop,
    "mp_reach_nlri": pack_mp_reach_nlri,
    "mp_unreach_nlri": pack_mp_unreach_nlri,
    "as4_path": pack_as4_path,
}

attribute_numbers = {
    "origin": 1,
    "as_path": 2,
    "next_hop": 3,
    "mp_reach_nlri": 14,
    "mp_unreach_nlri": 15,
    "as4_path": 17,
}

attribute_flags = {
    "origin": 0x40,
    "as_path": 0x40,
    "next_hop": 0x40,
    "mp_reach_nlri": 0x80,
    "mp_unreach_nlri": 0x80,
    "as4_path": 0xC0,
}


def parse_path_attributes(serialised_path_attributes, fourbyteas):
    stream = BytesIO(serialised_path_attributes)
    path_attributes = {}

    while True:
        attribute_header = stream.read(2)

        if len(attribute_header) == 0:
            break
        flags, type_code = struct.unpack("!BB", attribute_header)

        # TODO factor this pattern out
        if flags & 0x10:
            packed_length = stream.read(2)
            if len(packed_length) == 0:
                break
            (length,) = struct.unpack("!H", packed_length)
        else:
            packed_length = stream.read(1)
            if len(packed_length) == 0:
                break
            (length,) = struct.unpack("!B", packed_length)

        packed_attribute = stream.read(length)

        if type_code in attribute_parsers:
            # TODO we're ignoring the flags here, these should at very least be preserved
            # this is tightly coupled, there's gotta be a better way to do this
            if fourbyteas and attribute_keys[type_code] == "as_path":
                path_attributes["as_path"] = parse_as4_path(packed_attribute)
            else:
                path_attributes[attribute_keys[type_code]] = attribute_parsers[
                    type_code
                ](packed_attribute)
        else:
            print("WARNING did not recognise BGP path attribute type %d" % type_code)

    return path_attributes


def parse_withdrawn_routes(serialised_withdrawn_routes):
    stream = BytesIO(serialised_withdrawn_routes)
    prefixes = []

    while True:
        serialised_length = stream.read(1)
        if len(serialised_length) == 0:
            break
        prefix_length = ord(serialised_length)
        packed_prefix = stream.read(prefix_byte_length(prefix_length))
        prefix = unpack_prefix(packed_prefix)
        prefixes.append(IP4Prefix(prefix, prefix_length))

    return prefixes


PATH_ATTRIBUTE_ORDER = {
    "origin": 1,
    "as_path": 2,
    "as4_path": 3,
    "next_hop": 4,
    "mp_reach_nlri": 5,
    "mp_unreach_nlri": 6,
}


@register_parser
class BgpUpdateMessage(BgpMessage):
    MSG_TYPE = BgpMessage.UPDATE_MESSAGE

    def __init__(self, withdrawn_routes, path_attributes, nlri):
        self.withdrawn_routes = withdrawn_routes
        self.path_attributes = path_attributes
        self.nlri = nlri

    @classmethod
    def parse(cls, serialised_message, capabilities):
        data_stream = BytesIO(serialised_message)
        withdrawn_routes_length = bytes_to_short(data_stream.read(2))
        serialised_withdrawn_routes = data_stream.read(withdrawn_routes_length)
        withdrawn_routes = parse_withdrawn_routes(serialised_withdrawn_routes)

        total_path_attribute_length = bytes_to_short(data_stream.read(2))
        serialised_path_attributes = data_stream.read(total_path_attribute_length)
        path_attributes = parse_path_attributes(
            serialised_path_attributes, "fourbyteas" in capabilities
        )

        serialised_nlri = data_stream.read()
        nlri = parse_nlri(serialised_nlri)

        return cls(withdrawn_routes, path_attributes, nlri)

    def pack(self, capabilities):
        # TODO pack withdrawn routes
        packed_withdrawn_routes = self.pack_withdrawn_routes()
        packed_withdrawn_routes_length = struct.pack("!H", len(packed_withdrawn_routes))
        packed_path_attributes = self.pack_path_attributes("fourbyteas" in capabilities)
        packed_path_attributes_length = struct.pack("!H", len(packed_path_attributes))
        packed_nlri = self.pack_nlri()
        return (
            packed_withdrawn_routes_length
            + packed_withdrawn_routes
            + packed_path_attributes_length
            + packed_path_attributes
            + packed_nlri
        )

    def pack_withdrawn_routes(self):
        packed_routes_list = []

        for prefix in self.withdrawn_routes:
            packed_route = struct.pack("!B", prefix.length) + pack_prefix(
                prefix.prefix, prefix.length
            )
            packed_routes_list.append(packed_route)

        return b"".join(packed_routes_list)

    def pack_path_attributes(self, fourbyteas):
        packed_path_attributes = []
        sorted_attribute_pairs = sorted(
            self.path_attributes.items(), key=lambda x: PATH_ATTRIBUTE_ORDER[x[0]]
        )
        for name, path_attribute in sorted_attribute_pairs:
            if fourbyteas and name == "as_path":
                packed_entry = pack_as4_path(path_attribute)
            else:
                packed_entry = attribute_packers[name](path_attribute)
            packed_header = struct.pack(
                "!BBB",
                attribute_flags[name],
                attribute_numbers[name],
                len(packed_entry),
            )
            packed_path_attribute = packed_header + packed_entry
            packed_path_attributes.append(packed_path_attribute)

        return b"".join(packed_path_attributes)

    def pack_nlri(self):
        packed_nlri = []

        for prefix in self.nlri:
            # TODO this feels like it should be on IP4Prefix
            packed_prefix = struct.pack("!B", prefix.length) + pack_prefix(
                prefix.prefix, prefix.length
            )
            packed_nlri.append(packed_prefix)

        return b"".join(packed_nlri)

    def __str__(self):
        return (
            "BgpUpdateMessage: Withdrawn routes: %s, Path attributes: %s, NLRI: %s"
            % (
                [str(x) for x in self.withdrawn_routes],
                self.path_attributes,
                [str(x) for x in self.nlri],
            )
        )


@register_parser
class BgpNotificationMessage(BgpMessage):
    MSG_TYPE = BgpMessage.NOTIFICATION_MESSAGE

    MESSAGE_HEADER_ERROR = 1
    OPEN_MESSAGE_ERROR = 2
    UPDATE_MESSAGE_ERROR = 3
    HOLD_TIMER_EXPIRED = 4
    FINITE_STATE_MACHINE_ERROR = 5
    CEASE = 6

    def __init__(self, error_code, error_subcode=0, data=b""):
        self.error_code = error_code
        self.error_subcode = error_subcode
        self.data = data

    @classmethod
    def parse(cls, serialised_message, _capabilities):
        error_code, error_subcode = struct.unpack("!BB", serialised_message[:2])
        data = serialised_message[2:]
        return cls(error_code, error_subcode, data)

    def pack(self, _capabilities):
        return (
            struct.pack(
                "!BB",
                self.error_code,
                self.error_subcode,
            )
            + self.data
        )

    def __str__(self):
        return "BgpNotificationMessage: Error code: %s, Error subcode: %s, Data: %s" % (
            self.error_code,
            self.error_subcode,
            self.data,
        )


@register_parser
class BgpKeepaliveMessage(BgpMessage):
    MSG_TYPE = BgpMessage.KEEPALIVE_MESSAGE

    def __init__(self):
        pass

    @classmethod
    def parse(cls, serialised_message, _capabilities):
        return cls()

    def pack(self, _capabilities):
        return b""

    def __eq__(self, other):
        return type(other) == type(self)

    def __str__(self):
        return "BgpKeepaliveMessage"
