"""
firewall.py -
Module for interacting with the Mikrotik routeros shell to manage firewall rules and IP address lists.
Provides functions to read the current firewall configuration, add and remove IPs from the blocklist.
Reads the latest report from the IDS scan, reads the current firewall configuration, and updates the router's blocklist accordingly.
"""

from __future__ import annotations

import os
import json
import paramiko
from datetime import datetime
import shlex
from typing import Iterable
from .utils.constants import (
    BLOCKLIST_NAME,
    MIKROTIK_DATA_DIR,
    REPORT_DIR,
    ROUTER_IP,
    ROUTER_PASSWORD,
    ROUTER_USER,
    SSH_PORT,
    SSH_TIMEOUT,
    WHITELIST_NAMES,
)
from .utils.files import write_lines
from .utils.ip_utils import is_whitelisted
from .utils.blocklist_history import compute_changes, cumulative_blocklist


class MikrotikFirewall:
    """
    Class to manage Mikrotik firewall rules and IP address lists.
    Provides methods to read the current blocklist, add/remove IPs, and update the blocklist based on IDS scan reports.
    """

    def __init__(self, router_ip: str = ROUTER_IP, username: str = ROUTER_USER,
                 password: str = ROUTER_PASSWORD):
        self.router_ip = router_ip
        self.username = username
        self.password = password
        self.client: paramiko.SSHClient | None = None

    def connect(self) -> None:
        """Establish an SSH connection to the Mikrotik router."""
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            # RouterOS 7.20.x exposes an RSA host key and classic SSH ciphers.
            client.connect(
                self.router_ip,
                port=SSH_PORT,
                username=self.username,
                password=self.password,
                timeout=SSH_TIMEOUT,
                banner_timeout=SSH_TIMEOUT,
                auth_timeout=SSH_TIMEOUT,
                allow_agent=False,
                look_for_keys=False,
                compress=False,
                disabled_algorithms={"pubkeys": ["rsa-sha2-512"]},
            )
        except Exception:
            client.close()
            raise
        self.client = client

    def disconnect(self) -> None:
        """Close the SSH connection."""
        if self.client:
            self.client.close()
            self.client = None

    def execute_command(self, command: str) -> str:
        """Execute a command on the Mikrotik router via SSH and return the output."""
        if not self.client:
            raise RuntimeError("SSH client is not connected.")
        stdin, stdout, stderr = self.client.exec_command(command)
        output = stdout.read().decode("utf-8")
        error = stderr.read().decode("utf-8")
        if error:
            raise RuntimeError(f"Error executing command '{command}': {error}")
        return output

    def get_address_list_ipv4(self, name: str) -> set[str]:
        """Retrieve the current IPv4 address list from the Mikrotik router."""
        output = self.execute_command("/ip firewall address-list print")
        ips = set()
        for line in output.splitlines():
            parts = line.split()
            if len(parts) >= 3 and parts[1] == "list" and parts[2] == name:
                ip = parts[0]
                ips.add(ip)
        return ips

    def get_address_list_ipv6(self, name: str) -> set[str]:
        """Retrieve the current IPv6 address list from the Mikrotik router."""
        output = self.execute_command("/ipv6 firewall address-list print")
        ips = set()
        for line in output.splitlines():
            parts = line.split()
            if len(parts) >= 3 and parts[1] == "list" and parts[2] == name:
                ip = parts[0]
                ips.add(ip)
        return ips

    @staticmethod
    def _parse_terse_output(output: str) -> list[dict[str, str]]:
        """Parse RouterOS `print terse` records into key/value dictionaries."""
        records = []
        for line in output.splitlines():
            fields = {}
            for field in shlex.split(line):
                key, separator, value = field.partition("=")
                if separator:
                    fields[key] = value
            if fields:
                records.append(fields)
        return records

    def get_address_lists(self, ipv6: bool = False) -> list[dict[str, str]]:
        """Retrieve all address-list entries in structured form."""
        path = "/ipv6" if ipv6 else "/ip"
        output = self.execute_command(
            f"{path} firewall address-list print terse")
        return self._parse_terse_output(output)

    def get_firewall_rules(self, ipv6: bool = False) -> list[dict[str, str]]:
        """Retrieve all firewall filter rules in structured form."""
        path = "/ipv6" if ipv6 else "/ip"
        output = self.execute_command(
            f"{path} firewall filter print terse")
        return self._parse_terse_output(output)

    @staticmethod
    def _write_json(path: str, obj) -> None:
        with open(path, "w", encoding="utf-8") as file:
            json.dump(obj, file, indent=2, sort_keys=True)

    def save_state_report(self, report_dir: str = REPORT_DIR,
                          changes: dict | None = None) -> str:
        """Save current address lists and firewall rules for auditing.

        Writes structured snapshots to ``<report_dir>/mikrotik/data/`` and to
        the common ``data/mikrotik/`` directory (scratch mirror for debugging;
        **not** a git audit trail), plus a timestamped JSON report under
        ``<report_dir>/mikrotik/``. The authoritative long-term audit trail is
        the SQLite ``firewall_snapshots`` tables (see ``app.db``), which are
        appended here regardless of the file writes.
        """
        state = {
            "router": self.router_ip,
            "collected_at": datetime.now().astimezone().isoformat(),
            "address_lists": {
                "ipv4": self.get_address_lists(),
                "ipv6": self.get_address_lists(ipv6=True),
            },
            "firewall_rules": {
                "ipv4": self.get_firewall_rules(),
                "ipv6": self.get_firewall_rules(ipv6=True),
            },
        }
        if changes is not None:
            state["changes"] = changes

        report_data_dir = os.path.join(report_dir, "mikrotik", "data")
        os.makedirs(report_data_dir, exist_ok=True)
        os.makedirs(MIKROTIK_DATA_DIR, exist_ok=True)
        for version in ("ipv4", "ipv6"):
            for data_dir in (report_data_dir, MIKROTIK_DATA_DIR):
                self._write_json(
                    os.path.join(data_dir, f"address_list_{version}.json"),
                    state["address_lists"][version])
                self._write_json(
                    os.path.join(data_dir, f"firewall_{version}.json"),
                    state["firewall_rules"][version])

        # Persist a SQLite snapshot (authoritative audit trail, independent of
        # git). Best-effort: never breaks a router sync if the DB is down.
        self._persist_snapshot_to_db(state)

        timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
        report_path = os.path.join(report_dir, "mikrotik", f"{timestamp}.json")
        self._write_json(report_path, state)
        return report_path

    @staticmethod
    def _persist_snapshot_to_db(state: dict) -> None:
        """Write the firewall snapshot into the SQLite store (best-effort)."""
        try:
            from .db import get_connection
            from .db import repositories as repo
            repo.save_firewall_snapshot(
                get_connection(),
                address_lists_v4=state["address_lists"]["ipv4"],
                address_lists_v6=state["address_lists"]["ipv6"],
                rules_v4=state["firewall_rules"]["ipv4"],
                rules_v6=state["firewall_rules"]["ipv6"],
            )
        except Exception as exc:
            print(f"WARN: could not persist firewall snapshot to db ({exc})")

    def sync_blocklist(self, blocked_ipv4: Iterable[str],
                       blocked_ipv6: Iterable[str],
                       dry_run: bool = False,
                       report_dir: str = REPORT_DIR) -> dict[str, list[str]]:
        """Synchronize the Blocked list cumulatively and audit the result.

        The blocklist is grow-only: targets are added when present in the
        complete cumulative set (git history of the blocklist files + the
        current working files + this scan's targets) and never removed
        automatically. Only router entries absent from the complete set — i.e.
        never blocked by the system — are pruned. The history now lives in the
        SQLite blocklist ledger (``app.db`` ``blocklist_entries``); git history
        remains a read-only fallback. See ``app.blocklist_history``.
        """

        desired_ipv4 = set(blocked_ipv4)
        desired_ipv6 = set(blocked_ipv6)
        current_ipv4 = self.get_address_list_ipv4(BLOCKLIST_NAME)
        current_ipv6 = self.get_address_list_ipv6(BLOCKLIST_NAME)
        whitelisted_ipv4 = set()
        whitelisted_ipv6 = set()
        for name in WHITELIST_NAMES:
            whitelisted_ipv4.update(self.get_address_list_ipv4(name))
            whitelisted_ipv6.update(self.get_address_list_ipv6(name))

        v4_path = os.path.join(
            report_dir, "app", "data", "blocklist_ipv4.txt")
        v6_path = os.path.join(
            report_dir, "app", "data", "blocklist_ipv6.txt")
        complete_v4, complete_v6 = cumulative_blocklist(
            v4_path, v6_path, desired_ipv4, desired_ipv6)

        changes = compute_changes(
            current_ipv4, current_ipv6,
            complete_v4, complete_v6,
            whitelisted_ipv4, whitelisted_ipv6)
        changes["dry_run"] = dry_run

        if not dry_run:
            for ip in changes["removed_ipv4"]:
                self.execute_command(
                    f"/ip firewall address-list remove [find list={BLOCKLIST_NAME} address={ip}]")
            for ip in changes["removed_ipv6"]:
                self.execute_command(
                    f"/ipv6 firewall address-list remove [find list={BLOCKLIST_NAME} address={ip}]")
            self.write_address_list_ipv4(BLOCKLIST_NAME, changes["added_ipv4"])
            self.write_address_list_ipv6(BLOCKLIST_NAME, changes["added_ipv6"])
        self.save_state_report(report_dir, changes)
        return changes

    def write_address_list_ipv4(self, name: str, ips: Iterable[str]) -> None:
        """Write a list of IPv4 addresses to the Mikrotik router's address list."""
        for ip in ips:
            if is_whitelisted(ip):
                continue  # Skip whitelisted IPs
            self.execute_command(
                f"/ip firewall address-list add list={name} address={ip}")

    def write_address_list_ipv6(self, name: str, ips: Iterable[str]) -> None:
        """Write a list of IPv6 addresses to the Mikrotik router's address list."""
        for ip in ips:
            if is_whitelisted(ip):
                continue  # Skip whitelisted IPs
            self.execute_command(
                f"/ipv6 firewall address-list add list={name} address={ip}")


def write_blocklist_files(blocked_ipv4: Iterable[str],
                          blocked_ipv6: Iterable[str],
                          report_dir: str = REPORT_DIR) -> tuple[str, str]:
    """Write the final consolidated blocklist to ``<report_dir>/app/data/``.

    Returns the paths of the IPv4 and IPv6 blocklist files
    (``blocklist_ipv4.txt`` / ``blocklist_ipv6.txt``).
    """
    data_dir = os.path.join(report_dir, "app", "data")
    os.makedirs(data_dir, exist_ok=True)
    v4_path = os.path.join(data_dir, "blocklist_ipv4.txt")
    v6_path = os.path.join(data_dir, "blocklist_ipv6.txt")
    write_lines(v4_path, sorted(blocked_ipv4))
    write_lines(v6_path, sorted(blocked_ipv6))
    return v4_path, v6_path


def get_scan_report(report_dir: str = REPORT_DIR) -> dict[str, list[str]]:
    """Read the latest scan report `mikrotik-block.rsc` from the report directory and return the list of blocked IPs."""
    report_path = os.path.join(report_dir, "mikrotik-block.rsc")
    if not os.path.exists(report_path):
        raise FileNotFoundError(f"Scan report not found: {report_path}")

    # add sections, they start with `/ip firewall address-list` or `/ipv6 firewall address-list` and end with a blank line or EOF
    sections = {
        "/ip firewall address-list": [],
        "/ipv6 firewall address-list": [],
    }
    # add list=Cloudflare or list=Init7 to whitelist sections, list=Blocked to blocked section
    with open(report_path, "r", encoding="utf-8") as f:
        current_section = None
        for line in f:
            line = line.strip()
            if not line:
                current_section = None
                continue
            if line.startswith("/ip firewall address-list"):
                current_section = "/ip firewall address-list"
                continue
            if line.startswith("/ipv6 firewall address-list"):
                current_section = "/ipv6 firewall address-list"
                continue
            if current_section:
                sections[current_section].append(line)

    return sections


def main():
    """Legacy entry point for applying the latest generated RouterOS report."""
    # Read the latest scan report
    scan_report = get_scan_report(REPORT_DIR)

    # Connect to the Mikrotik router
    firewall = MikrotikFirewall()
    firewall.connect()

    # read current blocked IPs from the router
    current_blocked_ipv4 = firewall.get_address_list_ipv4(BLOCKLIST_NAME)
    current_blocked_ipv6 = firewall.get_address_list_ipv6(BLOCKLIST_NAME)

    # read whitelisted IPs from the router
    whitelisted_ipv4 = set()
    whitelisted_ipv6 = set()
    for wl_name in WHITELIST_NAMES:
        whitelisted_ipv4.update(firewall.get_address_list_ipv4(wl_name))
        whitelisted_ipv6.update(firewall.get_address_list_ipv6(wl_name))

    # read blocked IPs from the scan report
    blocked_ipv4 = set()
    blocked_ipv6 = set()
    for line in scan_report["/ip firewall address-list"]:
        if "list=Blocked" in line:
            parts = line.split()
            for part in parts:
                if part.startswith("address="):
                    ip = part.split("=")[1]
                    blocked_ipv4.add(ip)
    for line in scan_report["/ipv6 firewall address-list"]:
        if "list=Blocked" in line:
            parts = line.split()
            for part in parts:
                if part.startswith("address="):
                    ip = part.split("=")[1]
                    blocked_ipv6.add(ip)

    # Determine which IPs to add and remove
    to_add_ipv4 = blocked_ipv4 - current_blocked_ipv4 - whitelisted_ipv4
    to_add_ipv6 = blocked_ipv6 - current_blocked_ipv6 - whitelisted_ipv6
    to_remove_ipv4 = current_blocked_ipv4 - blocked_ipv4
    to_remove_ipv6 = current_blocked_ipv6 - blocked_ipv6

    # Remove old blocked IPs
    for ip in to_remove_ipv4:
        firewall.execute_command(
            f"/ip firewall address-list remove [find list={BLOCKLIST_NAME} address={ip}]")
    for ip in to_remove_ipv6:
        firewall.execute_command(
            f"/ipv6 firewall address-list remove [find list={BLOCKLIST_NAME} address={ip}]")

    # Add new blocked IPs
    firewall.write_address_list_ipv4(BLOCKLIST_NAME, to_add_ipv4)
    firewall.write_address_list_ipv6(BLOCKLIST_NAME, to_add_ipv6)

    firewall.save_state_report(REPORT_DIR)

    # Disconnect from the router
    firewall.disconnect()

    # Print summary
    print(
        f"Removed {len(to_remove_ipv4)} IPv4 and {len(to_remove_ipv6)} IPv6 addresses from the blocklist.")
    print(
        f"Added {len(to_add_ipv4)} IPv4 and {len(to_add_ipv6)} IPv6 addresses to the blocklist.")


if __name__ == "__main__":
    main()