"""Base classes and helpers for vulnerability evaluation.

Provides shared logic for reading CSVs and asserting conditions used across
mode-specific vulnerability checks.
"""
from __future__ import annotations

import csv
import os
import re
from collections import defaultdict
from typing import Iterable

from ptnetinspector.utils.csv_helpers import read_role_node_csv, sort_and_deduplicate_vul_csv
from ptnetinspector.utils.ip_utils import (
    check_ipv6_addresses_generated_from_prefix,
    has_additional_data,
    is_global_unicast_ipv6,
    is_valid_ipv4,
    is_valid_ipv6,
)
from ptnetinspector.utils.path import get_csv_path
from ptnetinspector.utils.vuln_catalog import load_vuln_catalog_by_test


class VulnerabilityBase:
    def __init__(
        self,
        interface: str,
        mode: list[str],
        ipver,
        smac: str,
        network: str,
        prefix_len: int,
        pref: int,
        dns: list[str],
        role_file: str | None = None,
        time_incoming_file: str | None = None,
        time_all_file: str | None = None,
        vulnerability_file: str | None = None,
        target_codes: set[str] | None = None,
        target_macs: set[str] | None = None,
    ) -> None:
        self.interface = interface
        self.mode = mode
        self.ipver = ipver
        self.smac = smac
        self.network = network
        self.prefix_len = prefix_len
        self.pref = pref
        self.dns = dns
        self.target_macs = {mac.upper() for mac in target_macs} if target_macs else None

        self.role_file = role_file or get_csv_path("role_node.csv")
        self.time_incoming_file = time_incoming_file or get_csv_path("time_incoming.csv")
        self.time_all_file = time_all_file or get_csv_path("time_all.csv")
        self.vulnerability_file = vulnerability_file or get_csv_path("vulnerability.csv")
        self.ra_file = get_csv_path("RA.csv")

        self.target_test_codes = {code.upper() for code in target_codes} if target_codes else None

    def _filter_by_target_macs(self) -> None:
        """Filter vulnerability CSV to include only target MACs (and Network row)."""
        if self.target_macs is None or not os.path.exists(self.vulnerability_file):
            return

        try:
            with open(self.vulnerability_file, "r", newline="") as csvfile:
                reader = csv.DictReader(csvfile)
                rows = []
                for row in reader:
                    entity_id = row.get("ID", "").strip()
                    mac = row.get("MAC", "").strip().upper()
                    # Keep Network rows (ID == "Network") and rows matching target MACs
                    if entity_id == "Network" or mac in self.target_macs:
                        rows.append(row)
        except Exception:
            return

        try:
            with open(self.vulnerability_file, "w", newline="") as csvfile:
                fieldnames = ["ID", "MAC", "Mode", "IPver", "Code", "Description", "Label"]
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(rows)
        except Exception:
            pass

    def _filter_modes(self, mode_string: str) -> str:
        modes = [m.strip() for m in mode_string.split(",")]
        active_modes = [m for m in modes if m in self.mode]
        return ",".join(active_modes) if active_modes else ""

    def _load_ra_prefixes(self) -> set[str]:
        prefixes: set[str] = set()
        if not os.path.exists(self.ra_file):
            return prefixes

        try:
            with open(self.ra_file, "r", newline="") as csvfile:
                reader = csv.DictReader(csvfile)
                for row in reader:
                    prefix = row.get("Prefix", "").strip()
                    if not prefix:
                        continue
                    if "/" not in prefix:
                        prefix = f"{prefix}/64"
                    prefixes.add(prefix)
        except Exception:
            pass

        return prefixes

    def _load_existing_rows(self) -> set[tuple[str, ...]]:
        existing_rows: set[tuple[str, ...]] = set()
        if os.path.exists(self.vulnerability_file):
            try:
                with open(self.vulnerability_file, "r", newline="") as csvfile:
                    reader = csv.DictReader(csvfile)
                    for row in reader:
                        row_tuple = tuple(
                            row[field].strip()
                            for field in [
                                "ID",
                                "MAC",
                                "Mode",
                                "IPver",
                                "Code",
                                "Description",
                                "Label",
                            ]
                        )
                        existing_rows.add(row_tuple)
            except Exception:
                pass
        return existing_rows

    def _prune_vulnerability_file(self) -> None:
        if self.target_test_codes is None or not os.path.exists(self.vulnerability_file):
            return

        try:
            test_catalog = load_vuln_catalog_by_test()
            code_to_test: dict[str, str] = {}
            for test_code, entries in test_catalog.items():
                for entry in entries:
                    vuln_code = entry.get("Code", "").strip().upper()
                    if vuln_code:
                        code_to_test[vuln_code] = test_code
        except Exception:
            return

        try:
            with open(self.vulnerability_file, "r", newline="") as csvfile:
                reader = csv.DictReader(csvfile)
                rows = []
                for row in reader:
                    vuln_code = row.get("Code", "").strip().upper()
                    test_code = code_to_test.get(vuln_code, "")
                    if test_code.upper() in self.target_test_codes:
                        rows.append(row)
        except Exception:
            return

        try:
            with open(self.vulnerability_file, "w", newline="") as csvfile:
                fieldnames = ["ID", "MAC", "Mode", "IPver", "Code", "Description", "Label"]
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(rows)
        except Exception:
            pass

    def _write_vulnerabilities(self, new_rows: Iterable[dict]) -> None:
        rows = list(new_rows)
        if not rows:
            return

        try:
            os.makedirs(os.path.dirname(self.vulnerability_file), exist_ok=True)
            write_header = not os.path.exists(self.vulnerability_file)
            with open(self.vulnerability_file, mode="a", newline="") as csvfile:
                fieldnames = ["ID", "MAC", "Mode", "IPver", "Code", "Description", "Label"]
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                if write_header:
                    writer.writeheader()
                writer.writerows(rows)
        except Exception:
            pass

    @staticmethod
    def _create_vuln_row(row_tuple: tuple[str, ...]) -> dict[str, str]:
        return dict(
            zip(["ID", "MAC", "Mode", "IPver", "Code", "Description", "Label"], row_tuple)
        )

    def _add_vuln_if_new(
        self,
        row_tuple: tuple[str, ...],
        existing_rows: set[tuple[str, ...]],
        new_rows: list[dict[str, str]],
    ) -> None:
        if row_tuple in existing_rows:
            return

        if self.target_test_codes is not None:
            vuln_code = row_tuple[4].strip().upper()
            try:
                test_catalog = load_vuln_catalog_by_test()
                code_to_test = {}
                for test_code, entries in test_catalog.items():
                    for entry in entries:
                        code = entry.get("Code", "").strip().upper()
                        if code:
                            code_to_test[code] = test_code
                test_code = code_to_test.get(vuln_code, "")
                if test_code.upper() not in self.target_test_codes:
                    return
            except Exception:
                pass

        new_rows.append(self._create_vuln_row(row_tuple))
        existing_rows.add(row_tuple)

    def _get_ipver_values(self) -> list[str]:
        ipver_values: list[str] = []
        if getattr(self.ipver, "ipv4", False):
            ipver_values.append("4")
        if getattr(self.ipver, "ipv6", False):
            ipver_values.append("6")
        return ipver_values

    def _extract_address(self, packet: str, addr_type: str) -> str:
        candidates = re.findall(rf"{addr_type}\s*=\s*([0-9A-Fa-f:.]+)", packet)
        for addr in candidates:
            if re.fullmatch(r"[0-9A-Fa-f]{2}(?::[0-9A-Fa-f]{2}){5}", addr):
                continue
            if is_valid_ipv4(addr) or is_valid_ipv6(addr):
                return addr
        return ""

    def _check_icmp6_outrange(self, packet: str, all_prefixes: set[str]) -> bool:
        if "ICMPv6" not in packet:
            return False

        src_ip = self._extract_address(packet, "src")
        if not src_ip or not is_valid_ipv6(src_ip):
            return False

        if all_prefixes:
            for prefix in all_prefixes:
                try:
                    if check_ipv6_addresses_generated_from_prefix(src_ip, prefix):
                        return False
                except Exception:
                    pass

        if is_global_unicast_ipv6(src_ip):
            return True

        return False

    def _process_query_response_vulnerability(
        self,
        query_check,
        response_check,
        role_dict,
        ipver_values,
        device_code,
        device_desc,
        net_code,
        net_desc,
        mode="a,a+",
        additional_query_check=None,
        additional_response_check=None,
    ) -> None:
        if not os.path.exists(self.time_all_file) or not has_additional_data(self.time_all_file):
            return

        mode = self._filter_modes(mode)
        if not mode:
            return

        existing_rows = self._load_existing_rows()
        new_rows: list[dict[str, str]] = []

        for ipver_value in ipver_values:
            query_found = False
            first_query_idx = -1
            response_macs = set()

            try:
                with open(self.time_all_file, "r", newline="") as csvfile:
                    reader = csv.DictReader(csvfile)
                    current_index = 0

                    for row in reader:
                        packet = row.get("packet", "")
                        mac = row.get("MAC", "").strip()

                        version_match = (ipver_value == "4" and "version = 4" in packet) or (
                            ipver_value == "6" and "version = 6" in packet
                        )
                        if not version_match:
                            current_index += 1
                            continue

                        if not query_found and query_check(packet, mac, self.smac):
                            if additional_query_check is None or additional_query_check(packet):
                                query_found = True
                                first_query_idx = current_index

                        if query_found and current_index > first_query_idx:
                            if response_check(packet, mac):
                                if additional_response_check is None or additional_response_check(packet):
                                    response_macs.add(mac)

                        current_index += 1

                if not query_found:
                    for device_id, mac in role_dict.items():
                        row = (str(device_id), mac, mode, ipver_value, device_code, device_desc, "2")
                        self._add_vuln_if_new(row, existing_rows, new_rows)

                    net_row = ("Network", "", mode, ipver_value, net_code, net_desc, "2")
                    self._add_vuln_if_new(net_row, existing_rows, new_rows)
                else:
                    for device_id, mac in role_dict.items():
                        label = "1" if mac in response_macs else "0"
                        row = (str(device_id), mac, mode, ipver_value, device_code, device_desc, label)
                        self._add_vuln_if_new(row, existing_rows, new_rows)

                    net_label = "1" if response_macs else "0"
                    net_row = ("Network", "", mode, ipver_value, net_code, net_desc, net_label)
                    self._add_vuln_if_new(net_row, existing_rows, new_rows)

            except (FileNotFoundError, KeyError):
                continue

        self._write_vulnerabilities(new_rows)

    @staticmethod
    def read_role_dict(path: str) -> dict:
        role_dict = read_role_node_csv(path)
        return role_dict if isinstance(role_dict, dict) else {}

    @staticmethod
    def sort_and_dedupe(path: str) -> None:
        sort_and_deduplicate_vul_csv(path)
