"""
Parses .yaml merge file, as generated by infrastructures upon collisions and
generates yet another infrastructure from it.

This can be used to merge hosts of one infrastructure or merge mutliple
infrastructures with this as the priority infrastructure.
"""

from typing import Any

import yaml

from scans2any.helpers.utils import is_ipv4
from scans2any.internal import Host, Infrastructure, Service, SortedSet, printer


def parse(
    filename: str = "MERGE_FILE.yaml",
) -> tuple[Infrastructure, list[dict] | None]:
    """
    Parses .yaml merge file and generates an Infrastructure from it.

    Parameters
    ----------
    filename : str, optional
        Filename of the .yaml merge file, by default "MERGE_FILE.yaml"

    Returns
    -------
    Infrastructure
        Information from merge file as an infrastructure object
    list[dict]
        ruleset for auto-merge as defined in merge file
    """

    printer.section("Parsing of Merge File")
    printer.status(f"Filename: {filename}")

    with open(filename) as merge_file:
        loaded_data: dict[str, Any] = yaml.safe_load(merge_file) or {}

    merge_data: dict[str, dict] = loaded_data.get("manual-merge") or {}
    auto_merge_data: list[dict] | None = loaded_data.get("auto-merge")
    custom_entries: dict[str, dict] = loaded_data.get("custom-entries") or {}

    infra = Infrastructure(identifier="Merge File")

    if __check_for_ambiguity(merge_data):
        printer.failure("Fix ambiguous entries in merge file.")
        exit(1)

    for hostname, info in merge_data.items():
        # Create new host
        if is_ipv4(hostname):
            new_host = Host(
                address=hostname,
                hostnames=SortedSet(),
                os=SortedSet(info["os"]) if "os" in info else SortedSet(),
            )
        else:
            new_host = Host(
                hostnames=SortedSet([hostname]),
                os=SortedSet(info["os"]) if "os" in info else SortedSet(),
            )

        # Fill with tcp services
        tcp_ports = info.get("tcp_ports", {})
        for port, service_info in tcp_ports.items():
            new_service = Service(
                port=port,
                protocol="tcp",
                service_names=SortedSet(service_info["service_names"])
                if "service_names" in service_info
                else SortedSet(),
                banners=SortedSet(service_info["banners"])
                if "banners" in service_info
                else SortedSet(),
            )
            new_host.add_service(new_service)

        # Fill with udp services
        udp_ports = info.get("udp_ports", {})
        for port, service_info in udp_ports.items():
            new_service = Service(
                port=port,
                protocol="udp",
                service_names=SortedSet(service_info["service_names"])
                if "service_names" in service_info
                else SortedSet(),
                banners=SortedSet(service_info["banners"])
                if "banners" in service_info
                else SortedSet(),
            )
            new_host.add_service(new_service)

        infra.add_host(new_host)

    for host, entries in custom_entries.items():
        if not __custom_entries_format_ok(entries):
            printer.warning(f"Custom entries for {host} in wrong format! Skipping..")
            continue
        new_host = Host(
            address=host,
            hostnames=SortedSet(entries.get("hostnames", [])),
            os=SortedSet([os] if (os := entries.get("os")) else []),
        )
        for service in entries.get("ports", []):
            port, proto = service["port"].split("/")
            new_service = Service(
                port=int(port),
                protocol=proto,
                service_names=SortedSet([service.get("service-name", "")]),
                banners=SortedSet(
                    [banner] if (banner := service.get("banner")) else []
                ),
            )
            new_host.add_service(new_service)

        infra.add_host(new_host)

    printer.success(f"Parsing of {len(infra.hosts)} hosts finished without errors")
    return infra, auto_merge_data


def __custom_entries_format_ok(entries: dict | None) -> bool:
    """
    Validate the format of custom entries.

    The entries must:
    - Be a non-None dictionary.
    - Contain the keys: "hostnames", "os", and "ports".
    - Each item in "ports" must be a dictionary containing:
        - "port": A string starting with a digit, followed by proto.
        - "service-name": A non-empty string.
        - "banner": A non-empty string.

    Args:
        entries (dict | None): The entries to validate.

    Returns:
        bool: True if the format is valid, False otherwise.
    """
    # Define the required top-level keys
    required_keys = {"ports"}

    # Check if entries is a dictionary and contains all required keys
    if not isinstance(entries, dict):
        return False
    if not required_keys.issubset(entries):
        return False

    # Validate each port entry
    ports = entries.get("ports", [])
    if not isinstance(ports, list):
        return False

    for port_entry in ports:
        if not isinstance(port_entry, dict):
            return False

        # Define required keys for each port entry
        port_required_keys = {"port"}
        if not port_required_keys.issubset(port_entry):
            return False

        port_value = port_entry.get("port")
        if not isinstance(port_value, str):
            return False

        # Extract the port number before the '/' if present
        port_num = port_value.split("/")[0]
        if not port_num.isdigit():
            return False

    return True


def __check_for_ambiguity(merge_data) -> bool:
    """
    Checks for ambuguity in merge file data, i. e. multiple service names or
    banners for the same port.

    Returns
    -------
    bool
        `True` if there are ambiguous entries, else `False`
    """

    is_ambiguous = False

    for hostname, info in merge_data.items():
        # Multiple OS's for hostname
        if "os" in info and len(info["os"]) > 1:
            printer.warning(f"Multiple OS's for host {hostname}")
            is_ambiguous = True

        tcp_ports = info.get("tcp_ports", {})
        udp_ports = info.get("udp_ports", {})

        for port, service_info in tcp_ports.items():
            # Multiple service names for tcp port
            if "service_name" in service_info and len(service_info["service_name"]) > 1:
                printer.warning(
                    f"Multiple service names for tcp port {port} of host {hostname}."
                )
                is_ambiguous = True
            # Multiple banners for tcp port
            if "banners" in service_info and len(service_info["banners"]) > 1:
                printer.warning(
                    f"Multiple banners for tcp port {port} of host {hostname}."
                )
                is_ambiguous = True

        for port, service_info in udp_ports.items():
            # Multiple service names for udp port
            if "service_name" in service_info and len(service_info["service_name"]) > 1:
                printer.warning(
                    f"Multiple service names for udp port {port} of host {hostname}."
                )
                is_ambiguous = True
            # Multiple banners for udp port
            if "banners" in service_info and len(service_info["banners"]) > 1:
                printer.warning(
                    f"Multiple banners for udp port {port} of host {hostname}."
                )
                is_ambiguous = True

    return is_ambiguous
