import json
import os
import subprocess
import time
from pathlib import Path

SERVERS_FILE = Path(os.environ.get("SERVERS_FILE", "/cache/servers.json"))

_fw_cache = {}

DEFAULT_SERVERS = [
    {"name": "ubuntu-web (.65)", "host": "192.168.0.65", "user": "rob", "key": "/root/.ssh/id_ed25519", "sudo": True},
    {"name": "ubuntu VPS (.2)", "host": "66.179.94.204", "user": "rob", "key": "/root/.ssh/id_ed25519", "sudo": True},
    {"name": "jetson-nano (.23)", "host": "192.168.0.23", "user": "rob", "key": "/root/.ssh/id_ed25519", "sudo": True},
]


def load_servers():
    if SERVERS_FILE.exists():
        try:
            servers = json.loads(SERVERS_FILE.read_text())
            if servers:
                return servers
        except Exception:
            pass
    return DEFAULT_SERVERS


def ssh_exec(server, cmd, timeout=30):
    # Local agent mode: run directly on this host (no SSH)
    if server.get("host") == "local":
        try:
            r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout)
            return r.returncode, r.stdout, r.stderr
        except Exception as e:
            return 1, "", str(e)
    host = server.get("host")
    user = server.get("user", "rob")
    key = server.get("key", "/root/.ssh/id_ed25519")
    ssh_cmd = [
        "ssh", "-i", key,
        "-o", "StrictHostKeyChecking=no",
        "-o", "UserKnownHostsFile=/dev/null",
        "-o", "ConnectTimeout=10",
        "-o", "BatchMode=yes",
        f"{user}@{host}", cmd,
    ]
    try:
        r = subprocess.run(ssh_cmd, capture_output=True, text=True, timeout=timeout)
        return r.returncode, r.stdout, r.stderr
    except Exception as e:
        return 1, "", str(e)


def _sudo(server):
    return "sudo " if server.get("sudo", True) else ""


def firewall_status(server):
    key = f"{server.get('name')}:{server.get('host')}"
    now = time.time()
    item = _fw_cache.get(key)
    if item and now - item[0] < 120:
        return item[1]
    s = _sudo(server)
    status = {
        "name": server.get("name", server.get("host")),
        "host": server.get("host"),
        "reachable": False,
        "country_block": False,
        "scanner_block": False,
        "scanner_ips": 0,
        "country_cidrs": 0,
        "error": None,
    }

    # single combined SSH call for all firewall checks
    cmd = (
        "echo ok; "
        f"{s}ipset list blocked_countries 2>/dev/null | grep -c /; "
        f"{s}iptables -C INPUT -m set --match-set blocked_countries src -j DROP 2>/dev/null && echo cb_active || echo cb_inactive; "
        f"{s}iptables -C INPUT -j SCANNER-BLOCK 2>/dev/null && echo sb_active || echo sb_inactive; "
        f"{s}iptables -S SCANNER-BLOCK 2>/dev/null | grep -c DROP"
    )
    rc, out, err = ssh_exec(server, cmd)
    if rc != 0 or "ok" not in out:
        status["error"] = (err or out or "unreachable").strip()
        return status
    status["reachable"] = True
    lines = [l.strip() for l in out.splitlines() if l.strip()]
    try:
        status["country_cidrs"] = int(lines[1]) if len(lines) > 1 else 0
    except (ValueError, IndexError):
        status["country_cidrs"] = 0
    if len(lines) > 2:
        status["country_block"] = lines[2] == "cb_active"
    if len(lines) > 3:
        status["scanner_block"] = lines[3] == "sb_active"
    if len(lines) > 4:
        try:
            status["scanner_ips"] = int(lines[4])
        except ValueError:
            pass

    _fw_cache[key] = (now, status)
    return status


def block_ip(server, ip):
    s = _sudo(server)
    # validate IP to avoid shell injection
    if not ip or any(c not in "0123456789." for c in ip):
        return {"error": "invalid IP"}
    rc, out, err = ssh_exec(server, f"{s}iptables -C SCANNER-BLOCK -s {ip} -j DROP 2>/dev/null || {s}iptables -A SCANNER-BLOCK -s {ip} -j DROP")
    if rc != 0:
        return {"error": (err or out).strip()}
    return {"status": "blocked", "ip": ip}


def unblock_ip(server, ip):
    s = _sudo(server)
    if not ip or any(c not in "0123456789." for c in ip):
        return {"error": "invalid IP"}
    rc, out, err = ssh_exec(server, f"{s}iptables -D SCANNER-BLOCK -s {ip} -j DROP 2>/dev/null")
    if rc != 0:
        return {"error": (err or out).strip()}
    return {"status": "unblocked", "ip": ip}


def toggle_country(server, enabled):
    s = _sudo(server)
    rule_in = f"{s}iptables -C INPUT -m set --match-set blocked_countries src -j DROP 2>/dev/null"
    rule_du = f"{s}iptables -C DOCKER-USER -m set --match-set blocked_countries src -j DROP 2>/dev/null"
    add_in = f"{s}iptables -I INPUT 1 -m set --match-set blocked_countries src -j DROP"
    add_du = f"{s}iptables -I DOCKER-USER 1 -m set --match-set blocked_countries src -j DROP"
    del_in = f"{s}iptables -D INPUT -m set --match-set blocked_countries src -j DROP 2>/dev/null"
    del_du = f"{s}iptables -D DOCKER-USER -m set --match-set blocked_countries src -j DROP 2>/dev/null"
    if enabled:
        cmd = f"({rule_in}) || {add_in}; ({rule_du}) || {add_du}"
    else:
        cmd = f"{del_in}; {del_du}"
    rc, out, err = ssh_exec(server, cmd)
    if rc != 0:
        return {"error": (err or out).strip()}
    return {"status": "enabled" if enabled else "disabled", "type": "country"}


def toggle_scanner(server, enabled):
    s = _sudo(server)
    rule_in = f"{s}iptables -C INPUT -j SCANNER-BLOCK 2>/dev/null"
    rule_du = f"{s}iptables -C DOCKER-USER -j SCANNER-BLOCK 2>/dev/null"
    add_in = f"{s}iptables -I INPUT 1 -j SCANNER-BLOCK"
    add_du = f"{s}iptables -I DOCKER-USER 1 -j SCANNER-BLOCK"
    del_in = f"{s}iptables -D INPUT -j SCANNER-BLOCK 2>/dev/null"
    del_du = f"{s}iptables -D DOCKER-USER -j SCANNER-BLOCK 2>/dev/null"
    if enabled:
        cmd = f"({rule_in}) || {add_in}; ({rule_du}) || {add_du}"
    else:
        cmd = f"{del_in}; {del_du}"
    rc, out, err = ssh_exec(server, cmd)
    if rc != 0:
        return {"error": (err or out).strip()}
    return {"status": "enabled" if enabled else "disabled", "type": "scanner"}