import type { BUN_SQLITE_WGUI_CLIENT_RULES } from "@/db/types/db";
import parsePortList from "@/src/utils/parse-port-list";
import validateIpv4 from "@/src/utils/validate-ipv4";
import validateIpv4Cidr from "@/src/utils/validate-ipv4-cidr";
import deriveIptablesChainNames from "./derive-iptables-chain-names";
import normalizeIptablesDestination from "./normalize-iptables-destination";
import validateClientRule from "./validate-client-rule";

type HostIptablesClient = {
    id?: number | "";
    wg_ip_address?: string;
    rules?: BUN_SQLITE_WGUI_CLIENT_RULES[];
};

type Params = {
    host_id: number;
    interface_name: string;
    target_interface: string;
    host_subnet?: string;
    clients?: HostIptablesClient[];
};

const SAFE_IFACE = /^[A-Za-z][A-Za-z0-9._-]*$/;
const MULTIPORT_LIMIT = 15;

function joinArgs(parts: (string | undefined)[]) {
    return parts.filter(Boolean).join(" ");
}

function protocolsForRule({
    protocol,
    has_ports,
}: {
    protocol?: string;
    has_ports: boolean;
}) {
    if (has_ports) {
        if (protocol == "udp") {
            return ["udp"];
        }

        if (protocol == "tcp") {
            return ["tcp"];
        }

        return ["tcp", "udp"];
    }

    if (protocol == "tcp" || protocol == "udp") {
        return [protocol];
    }

    return [undefined];
}

function portMatch({ ports }: { ports: string[] }) {
    if (!ports[0]) {
        return;
    }

    if (ports.length == 1) {
        return `--dport ${ports[0]}`;
    }

    return `-m multiport --dports ${ports.join(",")}`;
}

function chunkPorts({ ports }: { ports: string[] }) {
    if (ports.length <= MULTIPORT_LIMIT) {
        return [ports];
    }

    const chunks: string[][] = [];

    for (let i = 0; i < ports.length; i += MULTIPORT_LIMIT) {
        chunks.push(ports.slice(i, i + MULTIPORT_LIMIT));
    }

    return chunks;
}

function acceptLines({
    forward_chain,
    input_chain,
    source,
    protocol,
    ports,
    host_wg_ip,
}: {
    forward_chain: string;
    input_chain: string;
    source: string;
    protocol?: string;
    ports: string[];
    host_wg_ip?: string;
}) {
    const lines: string[] = [];
    const protocols = protocolsForRule({
        protocol,
        has_ports: Boolean(ports[0]),
    });
    const port_chunks = ports[0] ? chunkPorts({ ports }) : [[]];

    for (let p = 0; p < protocols.length; p++) {
        const proto = protocols[p];

        for (let c = 0; c < port_chunks.length; c++) {
            const chunk = port_chunks[c] || [];
            const port_match_str = proto
                ? portMatch({ ports: chunk })
                : undefined;

            lines.push(
                joinArgs([
                    `iptables -A ${forward_chain}`,
                    `-s ${source}`,
                    proto ? `-p ${proto}` : undefined,
                    port_match_str,
                    `-j ACCEPT`,
                ]),
            );

            lines.push(
                joinArgs([
                    `iptables -A ${input_chain}`,
                    `-s ${source}`,
                    host_wg_ip ? `-d ${host_wg_ip}` : undefined,
                    proto ? `-p ${proto}` : undefined,
                    port_match_str,
                    `-j ACCEPT`,
                ]),
            );
        }
    }

    return lines;
}

export default function buildHostIptablesScripts({
    host_id,
    interface_name,
    target_interface,
    host_subnet,
    clients,
}: Params) {
    if (!Number.isSafeInteger(host_id) || host_id < 0 || host_id > 999999999) {
        return {
            success: false,
            msg: `Invalid host id`,
        };
    }

    if (!SAFE_IFACE.test(interface_name)) {
        return {
            success: false,
            msg: `Invalid interface name`,
        };
    }

    if (!SAFE_IFACE.test(target_interface)) {
        return {
            success: false,
            msg: `Invalid target interface`,
        };
    }

    const scoped_subnet = host_subnet?.trim();

    if (scoped_subnet && !validateIpv4Cidr({ value: scoped_subnet })) {
        return {
            success: false,
            msg: `Invalid host subnet "${scoped_subnet}"`,
        };
    }

    const { forward, input } = deriveIptablesChainNames({ host_id });

    const masquerade_match = scoped_subnet ? `-s ${scoped_subnet} ` : ``;
    const accept_lines: string[] = [];
    const host_clients = clients || [];

    for (let i = 0; i < host_clients.length; i++) {
        const client = host_clients[i];

        if (!client) {
            continue;
        }

        const source_ip = client.wg_ip_address?.trim();

        if (!source_ip) {
            return {
                success: false,
                msg: `Client ${client.id || i} is missing a WireGuard IP`,
            };
        }

        if (!validateIpv4({ ip: source_ip })) {
            return {
                success: false,
                msg: `Client ${client.id || i} has an invalid WireGuard IP`,
            };
        }

        const source = `${source_ip}/32`;
        const rules = client.rules || [];

        for (let r = 0; r < rules.length; r++) {
            const rule = rules[r];

            if (!rule) {
                continue;
            }

            const valid = validateClientRule({ rule });

            if (!valid.success) {
                return {
                    success: false,
                    msg: `Client ${client.id || i}: ${valid.msg}`,
                };
            }

            if (rule.rule_type == "all") {
                accept_lines.push(
                    `iptables -A ${forward} -s ${source} -j ACCEPT`,
                    scoped_subnet
                        ? `iptables -A ${input} -s ${source} -d ${scoped_subnet} -j ACCEPT`
                        : `iptables -A ${input} -s ${source} -j ACCEPT`,
                );

                continue;
            }

            const destination = normalizeIptablesDestination({
                destination: rule.destination,
            });

            const parsed_ports = parsePortList({ ports: rule.ports });

            if (!parsed_ports.success) {
                return {
                    success: false,
                    msg: `Client ${client.id || i}: ${parsed_ports.msg}`,
                };
            }

            accept_lines.push(
                ...acceptLines({
                    forward_chain: forward,
                    input_chain: input,
                    source,
                    protocol: rule.protocol || "any",
                    ports: parsed_ports.ports,
                    host_wg_ip: scoped_subnet,
                }),
            );
        }
    }

    accept_lines.push(
        `# Prevent clients from forwarding traffic back into the WireGuard network`,
        `iptables -A ${forward} -o ${interface_name} -j DROP`,
    );

    const post_up = [
        `#!/bin/bash`,
        ``,
        `iptables -D INPUT -i ${interface_name} -j ACCEPT 2>/dev/null || true`,
        `iptables -D OUTPUT -o ${interface_name} -j ACCEPT 2>/dev/null || true`,
        `iptables -D FORWARD -i ${interface_name} -j ACCEPT 2>/dev/null || true`,
        `iptables -D FORWARD -o ${interface_name} -j ACCEPT 2>/dev/null || true`,
        ``,
        `iptables -N ${forward} 2>/dev/null || true`,
        `iptables -F ${forward}`,
        `iptables -N ${input} 2>/dev/null || true`,
        `iptables -F ${input}`,
        ``,
        `iptables -A ${forward} -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT`,
        `iptables -A ${input} -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT`,
        ``,
        `iptables -A ${input} -i ${interface_name} -p udp --dport 53 -j ACCEPT`,
        `iptables -A ${input} -i ${interface_name} -p tcp --dport 53 -j ACCEPT`,
        ``,
        ...accept_lines,
        ``,
        `iptables -A ${forward} -j DROP`,
        `iptables -A ${input} -j DROP`,
        ``,
        `iptables -C FORWARD -i ${interface_name} -j ${forward} 2>/dev/null || iptables -I FORWARD 1 -i ${interface_name} -j ${forward}`,
        `iptables -C FORWARD -o ${interface_name} -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT 2>/dev/null || iptables -I FORWARD 1 -o ${interface_name} -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT`,
        `iptables -C INPUT -i ${interface_name} -j ${input} 2>/dev/null || iptables -I INPUT 1 -i ${interface_name} -j ${input}`,
        ``,
        `iptables -t nat -C POSTROUTING ${masquerade_match} -j MASQUERADE 2>/dev/null || iptables -t nat -A POSTROUTING ${masquerade_match} -j MASQUERADE`,
        ``,
    ].join("\n");

    const post_down = [
        `#!/bin/bash`,
        ``,
        `iptables -D FORWARD -i ${interface_name} -j ${forward} 2>/dev/null || true`,
        `iptables -D FORWARD -o ${interface_name} -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT 2>/dev/null || true`,
        `iptables -D INPUT -i ${interface_name} -j ${input} 2>/dev/null || true`,
        `iptables -D INPUT -i ${interface_name} -j ACCEPT 2>/dev/null || true`,
        `iptables -D OUTPUT -o ${interface_name} -j ACCEPT 2>/dev/null || true`,
        `iptables -D FORWARD -i ${interface_name} -j ACCEPT 2>/dev/null || true`,
        `iptables -D FORWARD -o ${interface_name} -j ACCEPT 2>/dev/null || true`,
        `iptables -F ${forward} 2>/dev/null || true`,
        `iptables -X ${forward} 2>/dev/null || true`,
        `iptables -F ${input} 2>/dev/null || true`,
        `iptables -X ${input} 2>/dev/null || true`,
        `iptables -t nat -D POSTROUTING ${masquerade_match} -j MASQUERADE 2>/dev/null || true`,
        ``,
    ].join("\n");

    return {
        success: true,
        post_up,
        post_down,
    };
}