# Copyright (C) 2026 Linuxfabrik <info@linuxfabrik.ch>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# On Debian systems, the complete text of the GNU General Public License
# version 2 can be found in /usr/share/common-licenses/GPL-2.
#
# SPDX-License-Identifier: GPL-2.0-or-later

"""CompilerDriver_nft: nftables compilation orchestrator.

This is the main entry point for nftables compilation. The run() method
orchestrates: preprocessor -> NAT compilation -> policy compilation ->
routing compilation -> nft script assembly -> file writing.

Output is a Bash shell script wrapper that applies the compiled nft
rules via a heredoc piped to ``nft -f``.
"""

from __future__ import annotations

import io
import os
import socket
import textwrap
import uuid
from pathlib import Path
from typing import TYPE_CHECKING

import sqlalchemy

from firewallfabrik.compiler._base import CompilerStatus
from firewallfabrik.core.objects import (
    NAT,
    Firewall,
    Policy,
    Routing,
    RuleSet,
)
from firewallfabrik.driver._compiler_driver import CompilerDriver
from firewallfabrik.driver._jinja2_template import Jinja2Template
from firewallfabrik.platforms.nftables import __compiler_version__

if TYPE_CHECKING:
    import sqlalchemy.orm

    from firewallfabrik.core._database import DatabaseManager

AF_INET = socket.AF_INET
AF_INET6 = socket.AF_INET6


def _prepend(prefix: str, text: str) -> str:
    """Prepend a string to every non-empty line."""
    if not text:
        return text
    lines = text.split('\n')
    return '\n'.join(prefix + line if line else '' for line in lines)


class CompilerDriver_nft(CompilerDriver):
    """Orchestrates full nftables compilation.

    Creates and runs NAT and policy compilers, then assembles the
    output into a Bash shell script wrapper.
    """

    def __init__(self, db: DatabaseManager) -> None:
        super().__init__(db)
        self.have_nat: bool = False
        self.have_filter: bool = False

    def run(
        self,
        cluster_id: str,
        fw_id: str,
        single_rule_id: str,
    ) -> str:
        """Main compilation entry point.

        Performs the full nftables compilation pipeline:
        1. Look up firewall object
        2. Create OS configurator
        3. For each address family (IPv4/IPv6):
           a. Compile NAT rules
           b. Compile policy rules
        4. Compile routing rules
        5. Assemble shell script with embedded nft rules
        6. Write output file
        """
        from firewallfabrik.platforms.nftables._os_configurator import (
            OSConfigurator_nft,
        )

        # -- Look up firewall --
        with self.db.session() as session:
            if fw_id:
                fw_uuid = uuid.UUID(fw_id) if isinstance(fw_id, str) else fw_id
                fw = session.execute(
                    sqlalchemy.select(Firewall).where(
                        Firewall.id == fw_uuid,
                    ),
                ).scalar_one_or_none()
            else:
                self.error('No firewall ID provided')
                return ''

            if fw is None:
                self.error(f'Firewall {fw_id} not found')
                return ''

            self.fw = fw

            try:
                options = fw.options or {}

                self._warn_unsupported_options(options)

                # Create OS configurator
                oscnf = OSConfigurator_nft(session, fw)

                # Gather all rule sets
                all_policies = (
                    session.execute(
                        sqlalchemy.select(Policy).where(
                            Policy.device_id == fw.id,
                        ),
                    )
                    .scalars()
                    .all()
                )

                all_nat = (
                    session.execute(
                        sqlalchemy.select(NAT).where(
                            NAT.device_id == fw.id,
                        ),
                    )
                    .scalars()
                    .all()
                )

                # Determine whether to run IPv4/IPv6 compilation passes
                # based on the rule sets' explicit address-family flags.
                # If no rule set enables IPv6, skip the IPv6 pass entirely.
                self._any_rs_ipv6 = any(rs.ipv6 for rs in (*all_policies, *all_nat))
                if not self._any_rs_ipv6:
                    self.ipv6_run = False

                # Determine IPv4/IPv6 run order (based on GUI option)
                ipv4_6_runs: list[int] = []
                ipv4_6_order = options.get('ipv4_6_order', '')
                if not ipv4_6_order or ipv4_6_order == 'ipv4_first':
                    if self.ipv4_run:
                        ipv4_6_runs.append(AF_INET)
                    if self.ipv6_run:
                        ipv4_6_runs.append(AF_INET6)
                elif ipv4_6_order == 'ipv6_first':
                    if self.ipv6_run:
                        ipv4_6_runs.append(AF_INET6)
                    if self.ipv4_run:
                        ipv4_6_runs.append(AF_INET)

                # Collect all compiled rules per chain per AF
                # Structure: {chain_name: [rule_lines]}
                filter_chains: dict[str, list[str]] = {
                    'input': [],
                    'forward': [],
                    'output': [],
                }
                nat_chains: dict[str, list[str]] = {
                    'prerouting': [],
                    'postrouting': [],
                    'output': [],
                }

                for policy_af in ipv4_6_runs:
                    ipv6_policy = policy_af == AF_INET6

                    # Run preprocessor if we have rules
                    nat_count = sum(
                        1
                        for nat_rs in all_nat
                        if self._matching_address_family(nat_rs, policy_af)
                    )
                    policy_count = sum(
                        1
                        for pol_rs in all_policies
                        if self._matching_address_family(pol_rs, policy_af)
                    )

                    if nat_count or policy_count:
                        from firewallfabrik.platforms.linux._preprocessor import (
                            PreprocessorLinux,
                        )

                        prep = PreprocessorLinux(session, fw, ipv6_policy)
                        if single_rule_id:
                            prep.single_rule_compile_mode = True
                        prep.compile()

                    # --- NAT compilation ---
                    top_nat = None
                    for nat_rs in all_nat:
                        if not self._matching_address_family(nat_rs, policy_af):
                            continue
                        if self._is_top_ruleset(nat_rs):
                            top_nat = nat_rs
                            continue
                        self._process_nat_rule_set(
                            session,
                            fw,
                            nat_rs,
                            single_rule_id,
                            nat_chains,
                            oscnf,
                            policy_af,
                        )

                    if top_nat is not None:
                        self._process_nat_rule_set(
                            session,
                            fw,
                            top_nat,
                            single_rule_id,
                            nat_chains,
                            oscnf,
                            policy_af,
                        )

                    # --- Policy compilation ---
                    for all_top in range(2):
                        for pol_rs in all_policies:
                            if not self._matching_address_family(pol_rs, policy_af):
                                continue
                            is_top = self._is_top_ruleset(pol_rs)
                            if is_top and all_top == 0:
                                continue
                            if not is_top and all_top == 1:
                                continue

                            self._process_policy_rule_set(
                                session,
                                fw,
                                pol_rs,
                                single_rule_id,
                                filter_chains,
                                oscnf,
                                policy_af,
                            )

                # --- Routing compilation ---
                from firewallfabrik.platforms.linux._routing_compiler import (
                    RoutingCompilerLinux,
                )

                routing_output = ''
                routing_rs = (
                    session.execute(
                        sqlalchemy.select(Routing).where(
                            Routing.device_id == fw.id,
                        ),
                    )
                    .scalars()
                    .first()
                )

                if routing_rs:
                    routing_compiler = RoutingCompilerLinux(session, fw, False)
                    routing_compiler.set_source_ruleset(routing_rs)
                    routing_compiler.source_ruleset = routing_rs

                    if single_rule_id:
                        routing_compiler.single_rule_compile_mode = True
                        routing_compiler.single_rule_id = single_rule_id
                    routing_compiler.verbose = self.verbose > 0
                    routing_compiler.source_dir = self.source_dir

                    routing_rules_count = routing_compiler.prolog()
                    if routing_rules_count > 0:
                        routing_compiler.compile()
                        routing_compiler.epilog()

                    routing_output = routing_compiler.output.getvalue()

                    if routing_compiler.get_errors() or routing_compiler.get_warnings():
                        self.all_errors.extend(routing_compiler.get_errors())
                        self.all_errors.extend(routing_compiler.get_warnings())

                # --- Assemble nft rules body ---
                nft_rules_body = self._assemble_nft_rules_body(
                    fw,
                    oscnf,
                    filter_chains,
                    nat_chains,
                    self._any_rs_ipv6,
                )

                # Single-rule compile mode: return raw rules (no shell wrapper)
                if self.single_rule_compile_on:
                    errors_str = '\n'.join(self.all_errors)
                    return errors_str + nft_rules_body + routing_output

                # --- Determine output file names (needed for manifest) ---
                cluster_name = ''
                self.determine_output_file_names(fw, cluster_name)

                # --- Assemble shell script ---
                script = self._assemble_shell_script(
                    fw,
                    nft_rules_body,
                    routing_output,
                    oscnf,
                )

                # --- Write output file ---
                fw_id_str = str(fw.id)
                output_path = self.file_names.get(fw_id_str, '')
                if output_path:
                    self.info(f'Output file name: {output_path}')
                    try:
                        out_p = Path(output_path)
                        out_p.parent.mkdir(parents=True, exist_ok=True)
                        out_p.write_text(script, encoding='utf-8')
                        out_p.chmod(0o755)
                        if self.all_errors:
                            self.info(' Compiled with errors')
                        else:
                            self.info(' Compiled successfully')
                    except OSError as ex:
                        self.error(
                            f'Failed to open file {output_path} for writing: {ex}'
                        )
                        return str(ex)

            except Exception as ex:
                self._status = CompilerStatus.FWCOMPILER_ERROR
                return str(ex)

        return ''

    def _process_nat_rule_set(
        self,
        session: sqlalchemy.orm.Session,
        fw: Firewall,
        nat_rs: RuleSet,
        single_rule_id: str,
        nat_chains: dict[str, list[str]],
        oscnf,
        policy_af: int,
    ) -> None:
        """Compile a single NAT rule set."""
        from firewallfabrik.platforms.nftables._nat_compiler import (
            NATCompiler_nft,
        )

        ipv6_policy = policy_af == AF_INET6

        nat_compiler = NATCompiler_nft(session, fw, ipv6_policy, oscnf)
        nat_compiler.set_source_ruleset(nat_rs)
        nat_compiler.source_ruleset = nat_rs

        if single_rule_id:
            nat_compiler.single_rule_compile_mode = True
            nat_compiler.single_rule_id = single_rule_id
        nat_compiler.verbose = self.verbose > 0
        nat_compiler.source_dir = self.source_dir
        nat_compiler.debug_rule = self.debug_rule_nat
        nat_compiler.rule_debug_on = self.debug_rule_nat >= 0

        nat_rules_count = nat_compiler.prolog()
        if nat_rules_count > 0:
            nat_compiler.compile()
            nat_compiler.epilog()

        self.have_nat = self.have_nat or (nat_rules_count > 0)

        # Collect per-chain rules from the compiler
        for chain_name, rules in nat_compiler.chain_rules.items():
            if rules:
                nat_chains.setdefault(chain_name, []).extend(rules)

        if nat_compiler.get_errors() or nat_compiler.get_warnings():
            self.all_errors.extend(nat_compiler.get_errors())
            self.all_errors.extend(nat_compiler.get_warnings())

    def _process_policy_rule_set(
        self,
        session: sqlalchemy.orm.Session,
        fw: Firewall,
        pol_rs: RuleSet,
        single_rule_id: str,
        filter_chains: dict[str, list[str]],
        oscnf,
        policy_af: int,
    ) -> None:
        """Compile a single policy rule set."""
        from firewallfabrik.platforms.nftables._policy_compiler import (
            PolicyCompiler_nft,
        )

        ipv6_policy = policy_af == AF_INET6

        policy_compiler = PolicyCompiler_nft(session, fw, ipv6_policy, oscnf)
        policy_compiler.set_source_ruleset(pol_rs)
        policy_compiler.source_ruleset = pol_rs

        if single_rule_id:
            policy_compiler.single_rule_compile_mode = True
            policy_compiler.single_rule_id = single_rule_id
        policy_compiler.verbose = self.verbose > 0
        policy_compiler.source_dir = self.source_dir
        policy_compiler.debug_rule = self.debug_rule_policy
        policy_compiler.rule_debug_on = self.debug_rule_policy >= 0

        policy_rules_count = policy_compiler.prolog()
        if policy_rules_count > 0:
            policy_compiler.compile()
            policy_compiler.epilog()

        # Collect per-chain rules from the compiler
        for chain_name, rules in policy_compiler.chain_rules.items():
            if rules:
                self.have_filter = True
                filter_chains.setdefault(chain_name, []).extend(rules)

        if policy_compiler.get_errors() or policy_compiler.get_warnings():
            self.all_errors.extend(policy_compiler.get_errors())
            self.all_errors.extend(policy_compiler.get_warnings())

    def _assemble_nft_rules_body(
        self,
        fw: Firewall,
        oscnf,
        filter_chains: dict[str, list[str]],
        nat_chains: dict[str, list[str]],
        have_ipv6: bool,
    ) -> str:
        """Assemble the nft rules body for embedding in a heredoc.

        Returns only the nft-format content (table/chain declarations
        and inline rules) without any shebang, header comments, or
        routing commands.  Each table is flushed individually so that
        rules managed by other tools (Docker, CrowdSec, fail2ban)
        remain untouched.
        """
        options = fw.options or {}
        table_name = options.get('table_name', '') or 'fwf'
        filter_table = f'{table_name}_filter'
        nat_table = f'{table_name}_nat'

        out = io.StringIO()

        # Determine address family
        family = 'inet' if have_ipv6 else 'ip'
        nat_family = 'ip'

        # Default policy from firewall options
        input_policy = 'drop'
        output_policy = 'drop'
        forward_policy = 'drop'

        # --- Filter table ---
        input_rules = ''.join(filter_chains.get('input', []))
        forward_rules = ''.join(filter_chains.get('forward', []))
        output_rules = ''.join(filter_chains.get('output', []))
        have_filter = bool(
            input_rules.strip() or forward_rules.strip() or output_rules.strip()
        )

        # --- NAT table ---
        prerouting_rules = ''.join(nat_chains.get('prerouting', []))
        postrouting_rules = ''.join(nat_chains.get('postrouting', []))
        output_nat_rules = ''.join(nat_chains.get('output', []))
        have_nat = bool(
            prerouting_rules.strip()
            or postrouting_rules.strip()
            or output_nat_rules.strip()
        )

        # Atomically delete our tables before recreating them.
        # "create + delete" ensures deletion works even on first run
        # (plain "delete" fails if the table does not exist yet).
        if have_filter:
            out.write(f'table {family} {filter_table} {{}}\n')
            out.write(f'delete table {family} {filter_table}\n')
        if have_nat:
            out.write(f'table {nat_family} {nat_table} {{}}\n')
            out.write(f'delete table {nat_family} {nat_table}\n')
        if have_filter or have_nat:
            out.write('\n')

        if have_filter:
            out.write(f'table {family} {filter_table} {{\n')

            # Automatic rules (established/related, etc.)
            auto_rules = oscnf.generate_automatic_rules()

            # Input chain
            out.write('    chain input {\n')
            out.write(
                f'        type filter hook input priority filter; policy {input_policy};\n'
            )
            if auto_rules:
                out.write(auto_rules)
            if input_rules.strip():
                out.write(input_rules)
            out.write('    }\n')
            out.write('\n')

            # Forward chain
            out.write('    chain forward {\n')
            out.write(
                f'        type filter hook forward priority filter; policy {forward_policy};\n'
            )
            if auto_rules:
                out.write(auto_rules)
            if forward_rules.strip():
                out.write(forward_rules)
            out.write('    }\n')
            out.write('\n')

            # Output chain
            out.write('    chain output {\n')
            out.write(
                f'        type filter hook output priority filter; policy {output_policy};\n'
            )
            if auto_rules:
                out.write(auto_rules)
            if output_rules.strip():
                out.write(output_rules)
            out.write('    }\n')

            out.write('}\n')
            out.write('\n')

        if have_nat:
            out.write(f'table {nat_family} {nat_table} {{\n')

            # Prerouting chain (DNAT)
            out.write('    chain prerouting {\n')
            out.write('        type nat hook prerouting priority dstnat;\n')
            if prerouting_rules.strip():
                out.write(prerouting_rules)
            out.write('    }\n')

            # Output chain (local NAT — DNAT for locally-originated traffic)
            if output_nat_rules.strip():
                out.write('\n')
                out.write('    chain output {\n')
                out.write('        type nat hook output priority dstnat;\n')
                out.write(output_nat_rules)
                out.write('    }\n')

            out.write('\n')

            # Postrouting chain (SNAT/masquerade)
            out.write('    chain postrouting {\n')
            out.write('        type nat hook postrouting priority srcnat;\n')
            if postrouting_rules.strip():
                out.write(postrouting_rules)
            out.write('    }\n')

            out.write('}\n')
            out.write('\n')

        return out.getvalue()

    def _assemble_shell_script(
        self,
        fw: Firewall,
        nft_rules_body: str,
        routing_output: str,
        oscnf=None,
    ) -> str:
        """Assemble the complete shell script using the Jinja2 template."""
        options = fw.options or {}

        user_name = os.environ.get('USER', 'unknown')

        debug = options.get('debug', False)
        shell_debug = 'set -x' if debug else ''

        prolog_script = options.get('prolog_script', '')
        epilog_script = options.get('epilog_script', '')
        prolog_place = options.get('prolog_place', '') or 'top'

        nft_path = options.get('nft_path', '') or '/usr/sbin/nft'

        # Build comment block
        comment_text = (fw.comment or '').rstrip('\n')
        comment = _prepend('#  ', comment_text) if comment_text else ''

        # Build errors/warnings block
        errors_and_warnings = ''
        if self.all_errors:
            errors_and_warnings = '\n'.join(f'# {err}' for err in self.all_errors)

        # Management access for block action
        mgmt_access = bool(options.get('mgmt_access', False))
        ssh_management_address = options.get('mgmt_addr', '')

        # IP forwarding commands
        ip_forward_commands = self._get_ip_forward_commands(fw)

        # Interface configuration
        configure_interfaces = options.get('configure_interfaces', False)
        verify_interfaces_opt = options.get('verify_interfaces', False)
        ip_path = options.get('ip_path', '') or 'ip'

        shell_functions = ''
        configure_interfaces_code = ''
        verify_interfaces_code = ''

        if oscnf is not None:
            # Include shell functions when any interface feature is enabled
            need_shell_functions = (
                configure_interfaces
                or verify_interfaces_opt
                or options.get('configure_bridge_interfaces', False)
                or any(iface.is_dynamic() for iface in fw.interfaces)
            )
            if need_shell_functions:
                shell_functions = oscnf.print_shell_functions()

            if configure_interfaces:
                buf = io.StringIO()
                buf.write(oscnf.print_interface_configuration_commands())

                if options.get('configure_bridge_interfaces', False):
                    buf.write(oscnf.print_bridge_interface_configuration_commands())

                buf.write(oscnf.print_commands_to_clear_known_interfaces())
                buf.write(oscnf.print_dynamic_addresses_configuration_commands())
                raw = textwrap.dedent(buf.getvalue()).strip()
                configure_interfaces_code = textwrap.indent(raw, '    ')

            if verify_interfaces_opt:
                raw = textwrap.dedent(oscnf.print_verify_interfaces_commands()).strip()
                verify_interfaces_code = textwrap.indent(raw, '    ')

        # Build manifest line for installer
        fw_id_str = str(fw.id)
        local_name = Path(self.file_names.get(fw_id_str, '')).name
        manifest = f'# files: * {local_name}'
        remote = self.remote_file_names.get(fw_id_str, '')
        if remote:
            manifest += f' {remote}'

        # Named table support — only flush our own tables
        table_name = options.get('table_name', '') or 'fwf'
        filter_table = f'{table_name}_filter'
        nat_table = f'{table_name}_nat'

        # Determine filter family (must match _assemble_nft_rules_body)
        filter_family = 'inet' if self._any_rs_ipv6 else 'ip'

        context = {
            'version': __compiler_version__,
            'user': user_name,
            'comment': comment,
            'errors_and_warnings': errors_and_warnings,
            'manifest': manifest,
            'shell_debug': shell_debug,
            'nft_path': nft_path,
            'prolog_script': prolog_script,
            'epilog_script': epilog_script,
            'prolog_place': prolog_place,
            'nft_rules_body': nft_rules_body,
            'routing_output': routing_output,
            'ip_forward_commands': ip_forward_commands,
            'mgmt_access': mgmt_access,
            'ssh_management_address': ssh_management_address,
            'shell_functions': shell_functions,
            'configure_interfaces_code': configure_interfaces_code,
            'verify_interfaces_code': verify_interfaces_code,
            'filter_family': filter_family,
            'filter_table': filter_table,
            'flush_ruleset': options.get('flush_ruleset', True),
            'ip_path': ip_path,
            'nat_table': nat_table,
        }

        template = Jinja2Template('nftables', 'script.sh.j2')
        return template.render(context)

    def _get_ip_forward_commands(self, fw: Firewall) -> str:
        """Generate IP forwarding sysctl commands."""
        lines = []

        ipv4_fwd = str(fw.get_option('linux24_ip_forward') or '')
        if ipv4_fwd:
            val = 1 if ipv4_fwd in ('1', 'On', 'on') else 0
            lines.append(f'echo {val} > /proc/sys/net/ipv4/ip_forward')

        ipv6_fwd = str(fw.get_option('linux24_ipv6_forward') or '')
        if ipv6_fwd:
            val = 1 if ipv6_fwd in ('1', 'On', 'on') else 0
            lines.append(f'echo {val} > /proc/sys/net/ipv6/conf/all/forwarding')

        return '\n'.join(lines)

    # -- Utility methods --

    def _matching_address_family(self, ruleset: RuleSet, policy_af: int) -> bool:
        """Check if a rule set matches the given address family."""
        if hasattr(ruleset, 'matching_address_family'):
            return ruleset.matching_address_family(policy_af)
        return True

    def _is_top_ruleset(self, ruleset: RuleSet) -> bool:
        """Check if a rule set is the top-level rule set."""
        return bool(ruleset.top)

    def info(self, msg: str) -> None:
        """Print informational message."""
        if self.verbose:
            print(msg)
