"""Cascade (double-VPN) between two panel-managed WireGuard/AWG servers.

Traffic path:
  Clients → Entry server (accessible) → Exit server (blocked / target) → Internet

Implemented inside the entry Docker container as a second AWG/WG interface
(`cascade.conf`) with source-based routing for the VPN client subnet.
"""

from __future__ import annotations

import re
import shlex
from datetime import datetime
from typing import Any, Optional


WG_FAMILY = {'awg', 'awg2', 'awg_legacy', 'wireguard'}
CASCADE_MARKER = '# amnezia-web-panel-cascade'

DEFAULT_SETTINGS = {
    'pin_exit_route': True,          # host route to exit IP via real gateway
    'remove_eth_masquerade': True,   # stop NATing VPN clients out entry eth0/eth1
    'force_forward': True,           # FORWARD accept between awg0/wg0 and cascade
    'mss_clamp': True,               # fix broken TCP / endless page load
    'wait_handshake': True,          # fail if exit tunnel has no handshake
    'handshake_timeout_sec': 25,
    'allowed_ips': '0.0.0.0/0, ::/0',
    'keep_exit_dns': False,          # keep DNS= from exit client config
    'vpn_subnet_override': '',       # e.g. 10.8.1.0/24; empty = auto
    'table_id': 200,
    'rule_priority': 100,
}


def normalize_settings(raw: Optional[dict]) -> dict:
    settings = dict(DEFAULT_SETTINGS)
    if isinstance(raw, dict):
        for key, default in DEFAULT_SETTINGS.items():
            if key not in raw or raw[key] is None:
                continue
            if isinstance(default, bool):
                settings[key] = bool(raw[key])
            elif isinstance(default, int):
                try:
                    settings[key] = int(raw[key])
                except (TypeError, ValueError):
                    settings[key] = default
            else:
                settings[key] = str(raw[key]).strip()
    # Sanity clamps
    settings['handshake_timeout_sec'] = max(5, min(120, int(settings['handshake_timeout_sec'])))
    settings['table_id'] = max(1, min(252, int(settings['table_id'])))
    settings['rule_priority'] = max(1, min(32765, int(settings['rule_priority'])))
    if not settings.get('allowed_ips'):
        settings['allowed_ips'] = DEFAULT_SETTINGS['allowed_ips']
    return settings


class CascadeManager:
    def __init__(self, ssh_manager):
        self.ssh = ssh_manager

    @staticmethod
    def proto_base(protocol: str) -> str:
        return str(protocol or '').split('__', 1)[0]

    @staticmethod
    def is_wg_family(protocol: str) -> bool:
        return CascadeManager.proto_base(protocol) in WG_FAMILY

    @staticmethod
    def _container_name(protocol: str) -> str:
        base = CascadeManager.proto_base(protocol)
        match = re.search(r'__(\d+)$', str(protocol or ''))
        idx = int(match.group(1)) if match else 1
        names = {
            'awg': 'amnezia-awg',
            'awg2': 'amnezia-awg2',
            'awg_legacy': 'amnezia-awg-legacy',
            'wireguard': 'amnezia-wireguard',
        }
        name = names.get(base)
        if not name:
            return ''
        return name if idx <= 1 else f'{name}-{idx}'

    @staticmethod
    def _wg_binary(protocol: str) -> str:
        return 'wg' if CascadeManager.proto_base(protocol) in ('awg_legacy', 'wireguard') else 'awg'

    @staticmethod
    def _quick_binary(protocol: str) -> str:
        return 'wg-quick' if CascadeManager.proto_base(protocol) in ('awg_legacy', 'wireguard') else 'awg-quick'

    @staticmethod
    def _conf_dir(protocol: str) -> str:
        return '/opt/amnezia/wireguard' if CascadeManager.proto_base(protocol) == 'wireguard' else '/opt/amnezia/awg'

    @classmethod
    def _cascade_conf(cls, protocol: str) -> str:
        return f'{cls._conf_dir(protocol)}/cascade.conf'

    @classmethod
    def _cascade_up(cls, protocol: str) -> str:
        return f'{cls._conf_dir(protocol)}/cascade-up.sh'

    @staticmethod
    def _server_conf(protocol: str) -> str:
        base = CascadeManager.proto_base(protocol)
        if base == 'wireguard':
            return '/opt/amnezia/wireguard/wg0.conf'
        if base == 'awg_legacy':
            return '/opt/amnezia/awg/wg0.conf'
        return '/opt/amnezia/awg/awg0.conf'

    @staticmethod
    def prepare_exit_client_config(
        raw_config: str,
        endpoint_host: str,
        *,
        keep_dns: bool = False,
        allowed_ips: str = '0.0.0.0/0, ::/0',
    ) -> str:
        """Adapt a normal client config for use as an entry→exit cascade tunnel."""
        lines = []
        in_interface = False
        in_peer = False
        has_table = False
        for raw in (raw_config or '').splitlines():
            stripped = raw.strip()
            lower = stripped.lower()
            if stripped.startswith('[') and stripped.endswith(']'):
                in_interface = stripped.lower() == '[interface]'
                in_peer = stripped.lower() == '[peer]'
                lines.append(stripped)
                continue
            if in_interface and lower.startswith('dns') and not keep_dns:
                continue
            if in_interface and lower.startswith('table'):
                has_table = True
                lines.append('Table = off')
                continue
            if lower.startswith('endpoint'):
                port = '51820'
                if '=' in stripped:
                    rhs = stripped.split('=', 1)[1].strip()
                    if ':' in rhs:
                        port = rhs.rsplit(':', 1)[-1]
                host = (endpoint_host or '').strip()
                if host:
                    lines.append(f'Endpoint = {host}:{port}')
                else:
                    lines.append(stripped)
                continue
            if in_peer and lower.startswith('allowedips'):
                lines.append(f'AllowedIPs = {allowed_ips}')
                continue
            lines.append(stripped)

        if not has_table:
            out = []
            inserted = False
            for line in lines:
                out.append(line)
                if not inserted and line.strip().lower() == '[interface]':
                    out.append('Table = off')
                    inserted = True
            lines = out
        return '\n'.join(lines).strip() + '\n'

    @staticmethod
    def _extract_endpoint_host(config_text: str) -> str:
        for line in (config_text or '').splitlines():
            if line.strip().lower().startswith('endpoint'):
                rhs = line.split('=', 1)[-1].strip()
                return rhs.rsplit(':', 1)[0].strip().strip('[]')
        return ''

    def status(self, entry_protocol: str) -> dict:
        container = self._container_name(entry_protocol)
        if not container:
            return {'enabled': False, 'up': False, 'handshake': False, 'error': 'Unsupported protocol'}
        conf = self._cascade_conf(entry_protocol)
        wg_bin = self._wg_binary(entry_protocol)
        out, _, _ = self.ssh.run_sudo_command(
            f"docker exec {shlex.quote(container)} bash -lc "
            f"'echo HAS=$(test -f {shlex.quote(conf)} && echo 1 || echo 0); "
            f"echo IFACES=$({wg_bin} show interfaces 2>/dev/null); "
            f"echo HANDSHAKE=$({wg_bin} show cascade latest-handshakes 2>/dev/null | awk \"{{print \\$2}}\" | head -1); "
            f"echo TRANSFER=$({wg_bin} show cascade transfer 2>/dev/null | head -1); "
            f"ip rule show 2>/dev/null | head -20; "
            f"ip route show table 200 2>/dev/null | head -10'"
        )
        text = out or ''
        has_conf = 'HAS=1' in text
        ifaces_line = ''
        handshake_ts = 0
        for line in text.splitlines():
            if line.startswith('IFACES='):
                ifaces_line = line.split('=', 1)[1].strip()
            if line.startswith('HANDSHAKE='):
                try:
                    handshake_ts = int(line.split('=', 1)[1].strip() or '0')
                except ValueError:
                    handshake_ts = 0
        up = 'cascade' in ifaces_line.split()
        # WireGuard reports unix timestamp; 0 means never
        handshake_ok = handshake_ts > 0
        return {
            'enabled': has_conf,
            'up': up,
            'handshake': handshake_ok,
            'handshake_ts': handshake_ts,
            'raw': text.strip()[:3000],
        }

    def _vpn_subnet(self, entry_protocol: str, container: str, override: str = '') -> str:
        if override and re.match(r'^\d+\.\d+\.\d+\.\d+/\d+$', override.strip()):
            return override.strip()
        conf = self._server_conf(entry_protocol)
        out, _, _ = self.ssh.run_sudo_command(
            f"docker exec {shlex.quote(container)} bash -lc "
            f"\"grep -E '^Address' {shlex.quote(conf)} | head -1 | cut -d= -f2 | tr -d ' '\""
        )
        addr = (out or '').strip()
        if not addr:
            return '10.8.1.0/24'
        if '/' in addr:
            ip, cidr = addr.split('/', 1)
            parts = ip.split('.')
            if len(parts) == 4 and cidr.isdigit() and int(cidr) >= 24:
                return f"{parts[0]}.{parts[1]}.{parts[2]}.0/{cidr}"
            return addr
        parts = addr.split('.')
        if len(parts) == 4:
            return f"{parts[0]}.{parts[1]}.{parts[2]}.0/24"
        return '10.8.1.0/24'

    def apply(
        self,
        entry_protocol: str,
        exit_client_config: str,
        exit_host: str,
        settings: Optional[dict] = None,
    ) -> dict:
        opts = normalize_settings(settings)
        container = self._container_name(entry_protocol)
        if not container:
            return {'status': 'error', 'message': 'Unsupported entry protocol'}

        exists, _, code = self.ssh.run_sudo_command(
            f"docker inspect -f '{{{{.State.Running}}}}' {shlex.quote(container)} 2>/dev/null"
        )
        if code != 0 or exists.strip().lower() != 'true':
            return {'status': 'error', 'message': f'Entry container {container} is not running'}

        cascade_conf_path = self._cascade_conf(entry_protocol)
        cascade_up_path = self._cascade_up(entry_protocol)
        cascade_conf = self.prepare_exit_client_config(
            exit_client_config,
            exit_host,
            keep_dns=bool(opts['keep_exit_dns']),
            allowed_ips=str(opts['allowed_ips']),
        )
        endpoint_host = self._extract_endpoint_host(cascade_conf) or exit_host
        subnet = self._vpn_subnet(entry_protocol, container, opts.get('vpn_subnet_override') or '')
        wg_bin = self._wg_binary(entry_protocol)
        quick_bin = self._quick_binary(entry_protocol)
        conf_dir = self._conf_dir(entry_protocol)
        table_id = int(opts['table_id'])
        prio = int(opts['rule_priority'])
        hs_timeout = int(opts['handshake_timeout_sec'])

        pin_block = ''
        if opts['pin_exit_route']:
            pin_block = f"""
EXIT_IP="$EXIT_HOST"
if ! printf '%s' "$EXIT_IP" | grep -Eq '^[0-9]+\\.[0-9]+\\.[0-9]+\\.[0-9]+$'; then
  EXIT_IP=$(getent ahostsv4 "$EXIT_HOST" 2>/dev/null | awk '{{print $1; exit}}')
  if [ -z "$EXIT_IP" ]; then
    EXIT_IP=$(python3 - <<'PY' 2>/dev/null || true
import socket
print(socket.gethostbyname("{endpoint_host}"))
PY
)
  fi
fi
GW=$(ip route | awk '/default/ {{print $3; exit}}')
DEV=$(ip route | awk '/default/ {{print $5; exit}}')
if [ -n "$EXIT_IP" ] && [ -n "$GW" ] && [ -n "$DEV" ]; then
  ip route replace "$EXIT_IP/32" via "$GW" dev "$DEV" 2>/dev/null || true
fi
"""

        remove_masq = ''
        if opts['remove_eth_masquerade']:
            remove_masq = f"""
# Stop leaking VPN clients out of the entry server NIC
for ETH in eth0 eth1 eth2 ens3 ens5 enp0s3 enp1s0; do
  while iptables -t nat -D POSTROUTING -s "$SUBNET" -o "$ETH" -j MASQUERADE 2>/dev/null; do :; done
done
"""

        forward_block = ''
        if opts['force_forward']:
            forward_block = f"""
sysctl -w net.ipv4.ip_forward=1 >/dev/null 2>&1 || true
for SRC_IF in awg0 wg0; do
  iptables -C FORWARD -i "$SRC_IF" -o "$IFACE" -j ACCEPT 2>/dev/null \\
    || iptables -I FORWARD 1 -i "$SRC_IF" -o "$IFACE" -j ACCEPT 2>/dev/null || true
  iptables -C FORWARD -i "$IFACE" -o "$SRC_IF" -j ACCEPT 2>/dev/null \\
    || iptables -I FORWARD 1 -i "$IFACE" -o "$SRC_IF" -j ACCEPT 2>/dev/null || true
done
iptables -C FORWARD -i "$IFACE" -m state --state RELATED,ESTABLISHED -j ACCEPT 2>/dev/null \\
  || iptables -I FORWARD 1 -i "$IFACE" -m state --state RELATED,ESTABLISHED -j ACCEPT
"""

        mss_block = ''
        if opts['mss_clamp']:
            mss_block = """
iptables -t mangle -C FORWARD -p tcp --tcp-flags SYN,RST SYN -j TCPMSS --clamp-mss-to-pmtu 2>/dev/null \\
  || iptables -t mangle -A FORWARD -p tcp --tcp-flags SYN,RST SYN -j TCPMSS --clamp-mss-to-pmtu
"""

        handshake_block = ''
        if opts['wait_handshake']:
            handshake_block = f"""
OK=0
for i in $(seq 1 {hs_timeout}); do
  HS=$({shlex.quote(wg_bin)} show "$IFACE" latest-handshakes 2>/dev/null | awk '{{print $2}}' | head -1)
  if [ -n "$HS" ] && [ "$HS" != "0" ]; then OK=1; break; fi
  sleep 1
done
if [ "$OK" != "1" ]; then
  echo "CASCADE_NO_HANDSHAKE" >&2
  echo "Tunnel interface is up but exit peer did not handshake. Check exit server / UDP port / keys." >&2
  exit 42
fi
"""

        up_script = f"""#!/bin/bash
{CASCADE_MARKER}
set -e
CONF={shlex.quote(cascade_conf_path)}
QUICK={shlex.quote(quick_bin)}
SUBNET={shlex.quote(subnet)}
EXIT_HOST={shlex.quote(endpoint_host)}
IFACE=cascade
TABLE={table_id}
PRIO={prio}

"$QUICK" down "$CONF" 2>/dev/null || true
ip link delete "$IFACE" 2>/dev/null || true
ip rule del from "$SUBNET" table "$TABLE" 2>/dev/null || true
ip route flush table "$TABLE" 2>/dev/null || true

{pin_block}

"$QUICK" up "$CONF"

# Source-based routing: VPN clients go through cascade iface
ip route replace default dev "$IFACE" table "$TABLE" 2>/dev/null || ip route add default dev "$IFACE" table "$TABLE"
ip rule del from "$SUBNET" table "$TABLE" 2>/dev/null || true
ip rule add from "$SUBNET" table "$TABLE" priority "$PRIO"

{remove_masq}
{forward_block}
iptables -t nat -C POSTROUTING -s "$SUBNET" -o "$IFACE" -j MASQUERADE 2>/dev/null \\
  || iptables -t nat -A POSTROUTING -s "$SUBNET" -o "$IFACE" -j MASQUERADE
{mss_block}
{handshake_block}

echo CASCADE_UP_OK
"""

        self.ssh.run_sudo_command(f"docker exec {shlex.quote(container)} mkdir -p {shlex.quote(conf_dir)}")
        self.ssh.upload_file(cascade_conf, '/tmp/_amnz_cascade.conf')
        self.ssh.upload_file(up_script, '/tmp/_amnz_cascade_up.sh')
        self.ssh.run_sudo_command(
            f"docker cp /tmp/_amnz_cascade.conf {shlex.quote(container)}:{shlex.quote(cascade_conf_path)} && "
            f"docker cp /tmp/_amnz_cascade_up.sh {shlex.quote(container)}:{shlex.quote(cascade_up_path)} && "
            f"docker exec {shlex.quote(container)} chmod 644 {shlex.quote(cascade_conf_path)} && "
            f"docker exec {shlex.quote(container)} chmod +x {shlex.quote(cascade_up_path)}"
        )
        self.ssh.run_command('rm -f /tmp/_amnz_cascade.conf /tmp/_amnz_cascade_up.sh')

        self._ensure_start_hook(container, cascade_up_path)

        out, err, code = self.ssh.run_sudo_command(
            f"docker exec {shlex.quote(container)} bash {shlex.quote(cascade_up_path)}",
            timeout=max(90, hs_timeout + 40),
        )
        combined = ((out or '') + '\n' + (err or '')).strip()
        if code != 0 or 'CASCADE_UP_OK' not in (out or ''):
            msg = combined
            if 'CASCADE_NO_HANDSHAKE' in combined:
                msg = (
                    'Туннель к серверу выхода поднят, но handshake не прошёл. '
                    'Проверьте, что на выходе открыт UDP-порт, протокол совпадает, '
                    'и входной сервер может достучаться до IP выхода.'
                )
            return {
                'status': 'error',
                'message': (msg or 'Failed to bring cascade interface up').strip()[:900],
                'log': combined[:2000],
            }

        st = self.status(entry_protocol)
        return {
            'status': 'success',
            'subnet': subnet,
            'endpoint': endpoint_host,
            'container': container,
            'up': bool(st.get('up')),
            'handshake': bool(st.get('handshake')),
            'settings': opts,
            'applied_at': datetime.now().isoformat(timespec='seconds'),
            'diagnostics': st.get('raw', '')[:1500],
        }

    def _ensure_start_hook(self, container: str, cascade_up_path: str) -> None:
        marker = CASCADE_MARKER
        installer = f"""#!/bin/bash
set -e
START=/opt/amnezia/start.sh
[ -f "$START" ] || exit 0
grep -q '{marker}' "$START" 2>/dev/null && exit 0
printf '\\n%s\\n' '{marker}' >> "$START"
printf '%s\\n' 'if [ -x {cascade_up_path} ]; then bash {cascade_up_path} || true; fi' >> "$START"
"""
        self.ssh.upload_file(installer, '/tmp/_amnz_cascade_hook.sh')
        self.ssh.run_sudo_command(
            f"docker cp /tmp/_amnz_cascade_hook.sh {shlex.quote(container)}:/tmp/_hook.sh && "
            f"docker exec {shlex.quote(container)} bash /tmp/_hook.sh && "
            f"docker exec {shlex.quote(container)} rm -f /tmp/_hook.sh"
        )
        self.ssh.run_command('rm -f /tmp/_amnz_cascade_hook.sh')

    def disable(self, entry_protocol: str, settings: Optional[dict] = None) -> dict:
        opts = normalize_settings(settings)
        container = self._container_name(entry_protocol)
        if not container:
            return {'status': 'error', 'message': 'Unsupported entry protocol'}
        conf = self._cascade_conf(entry_protocol)
        up = self._cascade_up(entry_protocol)
        quick_bin = self._quick_binary(entry_protocol)
        table_id = int(opts['table_id'])
        script = f"""#!/bin/bash
set +e
QUICK={shlex.quote(quick_bin)}
CONF={shlex.quote(conf)}
UP={shlex.quote(up)}
TABLE={table_id}
"$QUICK" down "$CONF" 2>/dev/null
ip link delete cascade 2>/dev/null
ip rule del table "$TABLE" 2>/dev/null
ip rule del table "$TABLE" 2>/dev/null
ip route flush table "$TABLE" 2>/dev/null
rm -f "$CONF" "$UP"
echo CASCADE_DOWN_OK
"""
        self.ssh.upload_file(script, '/tmp/_amnz_cascade_down.sh')
        out, err, code = self.ssh.run_sudo_command(
            f"docker cp /tmp/_amnz_cascade_down.sh {shlex.quote(container)}:/tmp/_cascade_down.sh && "
            f"docker exec {shlex.quote(container)} bash /tmp/_cascade_down.sh && "
            f"docker exec {shlex.quote(container)} rm -f /tmp/_cascade_down.sh",
            timeout=60,
        )
        self.ssh.run_command('rm -f /tmp/_amnz_cascade_down.sh')
        if code != 0 and 'CASCADE_DOWN_OK' not in (out or ''):
            return {'status': 'success', 'message': (err or out or 'Cascade cleared (best effort)').strip()[:400]}
        return {'status': 'success'}