#!/usr/bin/env python3
"""MikroTik RouterOS script generation."""

from __future__ import annotations

import ipaddress
from datetime import datetime
from typing import Iterable

from .utils.ip_utils import is_ipv4, is_ipv6, is_whitelisted, parse_ip
from .utils.whitelists import WHITELISTS
from .utils.constants import BLOCKLIST_NAME


def split_cidr_versions(targets: Iterable[str]) -> tuple[list[str], list[str]]:
    """Split target CIDRs/IPs into (sorted IPv4 list, sorted IPv6 list)."""
    v4 = []
    v6 = []
    for t in targets:
        ip = t.split("/")[0]
        if is_ipv4(ip):
            v4.append(t)
        elif is_ipv6(ip):
            v6.append(t)

    def key_cidr(cidr: str):
        ip = cidr.split("/")[0]
        addr = parse_ip(ip)
        if addr is None:
            return (9, b"", cidr)
        return (addr.version, addr.packed, cidr)

    return sorted(v4, key=key_cidr), sorted(v6, key=key_cidr)


def mikrotik_script(tag: str, blocked_ips_or_subnets: Iterable[str],
                    whitelists: dict[str, list[str]] | None = None,
                    generated: str | None = None) -> tuple[str, int, int]:
    """Build a RouterOS .rsc script.

    Whitelisted ranges become `accept` rules placed above the `drop` rules.
    Blocked targets can be individual IPs (e.g. 1.2.3.4/32) or CIDR blocks
    (e.g. 1.2.3.0/24).

    Returns (script_text, blocked_count, whitelisted_cidr_count).
    """
    whitelists = WHITELISTS if whitelists is None else whitelists
    generated = generated or datetime.now().isoformat(timespec="seconds")

    def clean_blocked(targets: Iterable[str]) -> list[str]:
        cleaned = []
        for t in targets:
            ip = t.split("/")[0]
            if not is_whitelisted(ip):
                cleaned.append(t)
        return cleaned

    blocked = clean_blocked(set(blocked_ips_or_subnets))
    block_v4, block_v6 = split_cidr_versions(blocked)

    lines: list[str] = [f"# {tag} firewall export - generated {generated}", ""]
    wl_count = 0

    for name, cidrs in whitelists.items():
        nets = [ipaddress.ip_network(c) for c in cidrs]
        wl_count += len(nets)
        for version, cmd in ((4, "/ip firewall address-list"),
                             (6, "/ipv6 firewall address-list")):
            chunk = [n for n in nets if n.version == version]
            if not chunk:
                continue
            lines.append(f"# {name} IPv{version} ranges (whitelisted)")
            lines.append(cmd)
            lines += [f"add list={name} address={n}" for n in chunk]
            lines.append("")

    if block_v4:
        lines.append(f"# Blocked IPv4 ({len(block_v4)}, non-whitelisted)")
        lines.append("/ip firewall address-list")
        lines += [f"add list={BLOCKLIST_NAME} address={t}" for t in block_v4]
        lines.append("")
    if block_v6:
        lines.append(f"# Blocked IPv6 ({len(block_v6)}, non-whitelisted)")
        lines.append("/ipv6 firewall address-list")
        lines += [f"add list={BLOCKLIST_NAME} address={t}" for t in block_v6]
        lines.append("")

    for version, cmd, blocks in ((4, "/ip firewall filter", block_v4),
                                 (6, "/ipv6 firewall filter", block_v6)):
        lines.append(cmd)
        for name in whitelists:
            lines.append(
                f"add chain=input action=accept src-address-list={name} "
                f'comment="{tag} {name.lower()} whitelist (v{version})"')
        if blocks:
            for chain in ("input", "forward"):
                lines.append(
                    f"add chain={chain} action=drop src-address-list={BLOCKLIST_NAME} "
                    f'comment="{tag} blocked (v{version})"')
        lines.append("")

    return "\n".join(lines), len(blocked), wl_count


def write_mikrotik(path: str, tag: str,
                   blocked_ips_or_subnets: Iterable[str]) -> tuple[str, int, int]:
    text, n_blocked, n_whitelisted = mikrotik_script(
        tag, blocked_ips_or_subnets)
    with open(path, "w", encoding="utf-8") as f:
        f.write(text)
    return path, n_blocked, n_whitelisted