#!/usr/bin/python3
"""Own only table rpi_dualwan, route tables 20231/20232 and rule prefs 6000/6001."""
import concurrent.futures, json, pathlib, signal, subprocess, sys, threading, time
sys.path.insert(0, "/usr/local/lib")
from rpi_dualwan_adaptive import Adaptive
from rpi_dualwan_quality import eligible

PATH = '/usr/sbin:/usr/bin:/sbin:/bin'
UPLINKS = [('eth0', 20231, '0x11000000', 6000), ('wlan0', 20232, '0x12000000', 6001)]
stop = threading.Event()
reload_config = threading.Event()

def run(args, *, text=None, check=True):
    return subprocess.run(args, input=text, text=True, capture_output=True,
                          timeout=15, env={'PATH': PATH}, check=check)

def nft(text):
    run(['nft', '-f', '-'], text=text)

def cleanup():
    # AX9000 sends campus destinations here. Do not loop them back to AX9000 if CUP disappears.
    # Keep this reject route even when this service is stopped, until AX9000's campus route is removed.
    run(['ip', '-4', 'route', 'replace', 'unreachable', '10.0.0.0/8', 'proto', '186', 'metric', '2221'])
    run(['nft', 'delete', 'table', 'inet', 'rpi_dualwan'], check=False)
    for family in (4, 6):
        for iface, table, mark, pref in UPLINKS:
            while run(['ip', f'-{family}', 'rule', 'del', 'pref', str(pref),
                       'fwmark', mark + '/0x13000000', 'table', str(table)], check=False).returncode == 0:
                pass
            run(['ip', f'-{family}', 'route', 'flush', 'table', str(table)], check=False)
    run(['ip', '-4', 'route', 'del', '10.0.0.0/8', 'proto', '186', 'metric', '2220'], check=False)

def routes(family, iface):
    rows = json.loads(run(['ip', '-j', f'-{family}', 'route', 'show', 'default', 'dev', iface]).stdout)
    return min((r for r in rows if 'gateway' in r), key=lambda r:r.get('metric', 0), default=None)

def probe(family, iface, route):
    if not route:
        return False
    p = run(['curl', '--noproxy', '*', '--interface', iface, f'-{family}',
             '--head', '-s', '--connect-timeout', '3', '--max-time', '5',
             '-o', '/dev/null', '-w', '%{http_code}', 'https://www.baidu.com'], check=False)
    return p.returncode == 0 and p.stdout in ('200', '301', '302', '204')

def load_weights():
    result = {f'ipv{f}': {'eth0': 2, 'wlan0': 1} for f in (4, 6)}
    try:
        data = json.loads(pathlib.Path('/var/lib/rpi-dualwan/weights.json').read_text())
        for family in result:
            weights = data[family]['weights']
            if all(type(weights[i]) is int and 1 <= weights[i] <= 20 for i in ('eth0', 'wlan0')):
                result[family] = {i: weights[i] for i in ('eth0', 'wlan0')}
    except (OSError, ValueError, KeyError, TypeError):
        pass
    return result

def load_quality():
    try:
        data = json.loads(pathlib.Path('/var/lib/rpi-dualwan/weights.json').read_text())
        return data.get('quarantine', {}), data.get('quarantine_policy', {})
    except (OSError, ValueError):
        return {}, {}

def choices(family, healthy, weights):
    chain = f'choose{family}'
    prefix = f'add rule inet rpi_dualwan {chain} '
    result = [f'flush chain inet rpi_dualwan {chain}']
    if all(healthy):
        eth, cup = weights['eth0'], weights['wlan0']
        total = eth + cup
        buckets = ', '.join(f'{n} : jump select_eth' if ((n+1)*eth//total > n*eth//total)
                            else f'{n} : jump select_cup' for n in range(total))
        result.append(prefix + f'numgen inc mod {eth + cup} vmap {{ {buckets} }}')
    elif any(healthy):
        target = 'select_eth' if healthy[0] else 'select_cup'
        result.append(prefix + f'jump {target}')
    return result

def main():
    if len(sys.argv) > 1 and sys.argv[1] == 'stop':
        cleanup()
        return
    for sig in (signal.SIGTERM, signal.SIGINT):
        signal.signal(sig, lambda *_: stop.set())
    # One-time deployment handover: retain live routes, NAT, counters and conntrack.
    adopt = '--adopt' in sys.argv
    if adopt:
        run(['nft', 'list', 'table', 'inet', 'rpi_dualwan'])
    else:
        cleanup()
        run(['nft', '-f', '/etc/rpi-dualwan.nft'])
        for family in (4, 6):
            for iface, table, mark, pref in UPLINKS:
                run(['ip', f'-{family}', 'rule', 'add', 'pref', str(pref),
                     'fwmark', mark + '/0x13000000', 'table', str(table)])
    health = {(f,i[0]): False for f in (4,6) for i in UPLINKS}
    if adopt:
        old = json.loads(pathlib.Path('/run/rpi-dualwan.json').read_text())
        if time.time() - pathlib.Path('/run/rpi-dualwan.json').stat().st_mtime < 120:
            health = {(f,i): old.get(f'ipv{f}_{i}', False) is True for f,i in health}
    signal.signal(signal.SIGHUP, lambda *_: reload_config.set())
    adaptive = Adaptive(stop)
    streak = {key: [0,0] for key in health}
    previous_health = {f'ipv{f}_{i}': v for (f,i),v in health.items()} if adopt else None
    previous_effective = old.get('connection_weights') if adopt else None
    cached_routes = {}
    pending = {}
    next_health = next_decision = 0
    last_health = time.monotonic()
    next_error_log = 0
    info = {}
    state_path = pathlib.Path('/run/rpi-dualwan.json')
    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
        try:
            while not stop.is_set():
                now = time.monotonic()
                if reload_config.is_set():
                    # Validate in a temporary object before replacing the live reader.
                    temp_stop = threading.Event()
                    try:
                        replacement = Adaptive(temp_stop)
                        old = adaptive
                        old.stop = threading.Event()
                        old.stop.set()
                        old.close()
                        adaptive = replacement
                        print('adaptive configuration reloaded', flush=True)
                    except Exception as exc:
                        temp_stop.set()
                        print('adaptive reload rejected: ' + str(exc), flush=True)
                    reload_config.clear()
                    next_decision = 0
                if now >= next_health and not pending:
                    for family in (4,6):
                        for iface, table, mark, pref in UPLINKS:
                            key = (family, iface)
                            route = routes(family, iface)
                            gateway = route['gateway'] if route else None
                            # NM may delete per-interface routes during reapply even
                            # when the gateway address is unchanged. Restore defaults
                            # every health cycle without changing conntrack or marks.
                            if gateway != cached_routes.get(key) or gateway is not None:
                                if gateway:
                                    run(['ip', f'-{family}', 'route', 'replace', 'default', 'via', gateway,
                                         'dev', iface, 'onlink', 'table', str(table)])
                                else:
                                    run(['ip', f'-{family}', 'route', 'flush', 'table', str(table)], check=False)
                                cached_routes[key] = gateway
                            if family == 4 and iface == 'wlan0':
                                if gateway:
                                    run(['ip', '-4', 'route', 'replace', '10.0.0.0/8', 'via', gateway,
                                         'dev', iface, 'proto', '186', 'metric', '2220'])
                                else:
                                    run(['ip', '-4', 'route', 'del', '10.0.0.0/8', 'proto', '186',
                                         'metric', '2220'], check=False)
                            if not route:
                                health[key] = False
                            pending[key] = pool.submit(probe, family, iface, route)
                    next_health = now + 30
                if pending and all(f.done() for f in pending.values()):
                    for key, future in pending.items():
                        try:
                            ok = future.result()
                        except Exception:
                            ok = False
                        good, bad = streak[key]
                        streak[key] = [good + 1, 0] if ok else [0, bad + 1]
                        if streak[key][0] >= 2:
                            health[key] = True
                        elif streak[key][1] >= 2:
                            health[key] = False
                    pending = {}
                    last_health = now
                    next_decision = 0
                try:
                    adaptive.sample()
                except Exception as exc:
                    if now >= next_error_log:
                        print('adaptive sample error: ' + str(exc), flush=True)
                        next_error_log = now + 30
                if now >= next_decision:
                    base = load_weights()
                    quality, quality_policy = load_quality()
                    try:
                        starting = {name: (value if sum(value.values()) else base[name])
                                    for name, value in (previous_effective or base).items()}
                        weights, info = adaptive.decide({k:dict(v) for k,v in base.items()}, starting)
                    except Exception as exc:
                        weights, info = base, {'mode': 'baseline-error', 'error': str(exc)}
                    if now-last_health > 90:
                        health = dict.fromkeys(health, False)
                    current_health = {f'ipv{f}_{i}': v for (f,i),v in health.items()}
                    commands = ['flush set inet rpi_dualwan live']
                    effective = {}
                    emergencies = {}
                    for family in (4,6):
                        name = f'ipv{family}'
                        healthy, fallback = eligible([health[(family,i[0])] for i in UPLINKS],
                                                     quality.get(name, {}))
                        emergencies[name] = fallback
                        effective[name] = weights[name] if all(healthy) else {
                            'eth0': 20 if healthy[0] else 0, 'wlan0': 20 if healthy[1] else 0}
                        if (previous_health is None or previous_effective[name] != effective[name] or
                            any(previous_health[f'ipv{family}_{i[0]}'] != health[(family,i[0])] for i in UPLINKS)):
                            commands += choices(family, healthy, weights[name])
                        if any(healthy):
                            commands += [f'add element inet rpi_dualwan live {{ {name} timeout 120s }}']
                    nft('\n'.join(commands) + '\n')
                    if current_health != previous_health:
                        print('uplink health ' + json.dumps(current_health), flush=True)
                    if effective != previous_effective:
                        print('connection weights ' + json.dumps(effective) + ' mode=' + info['mode'], flush=True)
                    previous_health, previous_effective = current_health, effective
                    state = {**current_health, 'updated': time.strftime('%Y-%m-%dT%H:%M:%S%z'),
                             'connection_weights': effective, 'baseline_weights': base, 'adaptive': info,
                             'speed_quarantine': quality, 'quarantine_policy': quality_policy,
                             'quarantine_fallback': emergencies}
                    state_path.with_suffix('.tmp').write_text(json.dumps(state, indent=2) + '\n')
                    state_path.with_suffix('.tmp').replace(state_path)
                    next_decision = now + adaptive.config['decision_seconds']
                stop.wait(adaptive.config['sample_seconds'])
        finally:
            adaptive.stop.set()
            adaptive.close()

if __name__ == '__main__':
    main()
