"""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


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


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) -> str:
        """Adapt a normal client config for use as an entry→exit cascade tunnel."""
        lines = []
        in_interface = False
        has_table = False
        for raw in (raw_config or '').splitlines():
            line = raw.rstrip('\n')
            stripped = line.strip()
            lower = stripped.lower()
            if stripped.startswith('[') and stripped.endswith(']'):
                in_interface = stripped.lower() == '[interface]'
                lines.append(stripped)
                continue
            if in_interface and lower.startswith('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
            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, '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"'test -f {shlex.quote(conf)} && echo HAS_CONF; "
            f"ip -o link show cascade 2>/dev/null | head -1; "
            f"{wg_bin} show interfaces 2>/dev/null'"
        )
        text = out or ''
        ifaces = []
        for line in text.splitlines():
            if line.strip() == 'HAS_CONF':
                continue
            ifaces.extend(line.split())
        up = 'cascade' in ifaces or bool(re.search(r'\bcascade\b', text))
        return {
            'enabled': 'HAS_CONF' in text,
            'up': up,
            'raw': text.strip()[:2000],
        }

    def _vpn_subnet(self, entry_protocol: str, container: str) -> str:
        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) -> dict:
        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)
        endpoint_host = self._extract_endpoint_host(cascade_conf) or exit_host
        subnet = self._vpn_subnet(entry_protocol, container)
        wg_bin = self._wg_binary(entry_protocol)
        quick_bin = self._quick_binary(entry_protocol)
        conf_dir = self._conf_dir(entry_protocol)

        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

"$QUICK" down "$CONF" 2>/dev/null || true
ip link delete "$IFACE" 2>/dev/null || true

# Resolve hostname to IP if needed (pin route must use IP)
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}}')
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

"$QUICK" up "$CONF"

ip route replace default dev "$IFACE" table 200 2>/dev/null || ip route add default dev "$IFACE" table 200
ip rule del from "$SUBNET" table 200 2>/dev/null || true
ip rule add from "$SUBNET" table 200 priority 100

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
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
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

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=90,
        )
        if code != 0 or 'CASCADE_UP_OK' not in (out or ''):
            return {
                'status': 'error',
                'message': (err or out or 'Failed to bring cascade interface up').strip()[:800],
            }

        st = self.status(entry_protocol)
        return {
            'status': 'success',
            'subnet': subnet,
            'endpoint': endpoint_host,
            'container': container,
            'up': bool(st.get('up')),
            'applied_at': datetime.now().isoformat(timespec='seconds'),
        }

    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) -> dict:
        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)
        # Avoid nested quotes issues: upload a small script
        script = f"""#!/bin/bash
set +e
QUICK={shlex.quote(quick_bin)}
CONF={shlex.quote(conf)}
UP={shlex.quote(up)}
"$QUICK" down "$CONF" 2>/dev/null
ip link delete cascade 2>/dev/null
ip rule del table 200 2>/dev/null
ip rule del table 200 2>/dev/null
ip route flush table 200 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'}