#!/usr/bin/python3
"""Bounded public-mirror samples; atomically publish per-family connection weights."""
import fcntl
import ipaddress
import json
import math
import os
import pathlib
import re
import socket
import statistics
import subprocess
import sys
import tempfile
import time
from urllib.parse import urlsplit
sys.path.insert(0, '/usr/local/lib')
from rpi_dualwan_quality import policy, update as update_quality
from rpi_dualwan_game_guard import game_status, MESSAGES
from rpi_dualwan_capacity import archive, reserve_full_test

PATH = '/usr/sbin:/usr/bin:/sbin:/bin'
CONFIG = pathlib.Path('/etc/rpi-dualwan-speedtest.json')
STATE = pathlib.Path('/var/lib/rpi-dualwan/weights.json')
HEALTH = pathlib.Path('/run/rpi-dualwan.json')
FLOWS = pathlib.Path('/run/rpi-flow-monitor/status.json')
IFACES = ('eth0', 'wlan0')


def log(message):
    print(message, flush=True)


def read_state(path):
    try:
        data = json.loads(path.read_text())
        return data if isinstance(data, dict) else {}
    except (OSError, ValueError):
        return {}


class GameActivity(RuntimeError):
    """Skip the entire run without publishing partial measurements."""


def game_guard(config):
    state = game_status(read_state(FLOWS), config)
    if state['paused']:
        raise GameActivity(MESSAGES[state['reason']])


def run_download(args, timeout, config):
    game_guard(config)
    with subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                          text=True, env={'PATH': PATH}) as proc:
        deadline = time.monotonic() + timeout
        try:
            while True:
                game_guard(config)
                remaining = deadline-time.monotonic()
                if remaining <= 0:
                    raise subprocess.TimeoutExpired(args, timeout)
                try:
                    stdout, stderr = proc.communicate(timeout=min(1, remaining))
                    game_guard(config)
                    return subprocess.CompletedProcess(args, proc.returncode, stdout, stderr)
                except subprocess.TimeoutExpired:
                    if time.monotonic() >= deadline:
                        raise
        finally:
            # Only terminate this sampler's curl; no connection/routing services are touched.
            if proc.poll() is None:
                proc.terminate()
                try:
                    proc.communicate(timeout=1)
                except subprocess.TimeoutExpired:
                    proc.kill()
                    proc.communicate()


def smooth(old, shares, alpha, deadband, max_step):
    weights = old.get('weights', {'eth0': 2, 'wlan0': 1})
    if not all(type(weights.get(i)) is int and 1 <= weights[i] <= 20 for i in IFACES):
        weights = {'eth0': 2, 'wlan0': 1}
    applied = weights['eth0'] / sum(weights[i] for i in IFACES)
    previous = old.get('smoothed_eth_share', applied)
    if not isinstance(previous, (int, float)) or not math.isfinite(previous) or not 0 <= previous <= 1:
        previous = applied
    sample = statistics.median(shares)
    ema = (1 - alpha) * previous + alpha * sample
    if abs(ema - applied) >= deadband:
        target = max(0.1, min(0.9, ema))
        low = max(2, math.ceil((applied - max_step) * 20 - 1e-9))
        high = min(18, math.floor((applied + max_step) * 20 + 1e-9))
        slots = max(low, min(high, round(target * 20)))
        divisor = math.gcd(slots, 20 - slots)
        weights = {'eth0': slots // divisor, 'wlan0': (20 - slots) // divisor}
    return {'weights': weights, 'sample_eth_share': sample, 'smoothed_eth_share': ema}


def healthy(family, iface=None):
    state = read_state(HEALTH)
    try:
        fresh = time.time() - HEALTH.stat().st_mtime < 120
    except OSError:
        return False
    return fresh and all(state.get(f'ipv{family}_{i}') is True for i in (IFACES if iface is None else (iface,)))


def addresses(host, family):
    rows = socket.getaddrinfo(host, 443, socket.AF_INET if family == 4 else socket.AF_INET6,
                              socket.SOCK_STREAM)
    return sorted({r[4][0] for r in rows if ipaddress.ip_address(r[4][0]).is_global
                   and not ipaddress.ip_address(r[4][0]).is_multicast})


def measure(provider, family, iface, address, config):
    url = provider['url']
    host = urlsplit(url).hostname
    resolved = f'[{address}]' if family == 6 else address
    limit = config['max_bytes_per_sample']
    duration = config['duration_seconds']
    with tempfile.TemporaryDirectory(prefix='rpi-dualwan-test-', dir='/run') as temp:
        headers = pathlib.Path(temp) / 'headers'
        args = ['curl', '-q', '--noproxy', '*', f'-{family}', '--interface', iface,
                '--proto', '=https',
                '--resolve', f'{host}:443:{resolved}', '--fail', '--silent', '--show-error',
                '--connect-timeout', '3', '--max-time', str(duration),
                '--range', f'0-{limit-1}', '--max-filesize', str(limit),
                '--dump-header', str(headers), '--output', '/dev/null', '--write-out', '%{json}', url]
        p = run_download(args, duration + 5, config)
        data = json.loads(p.stdout or '{}')
        header = headers.read_text() if headers.exists() else ''
    size = data.get('size_download', 0)
    speed = data.get('speed_download', 0)
    elapsed = data.get('time_total', 0)
    response_ok = (p.returncode in (0, 28) and data.get('http_code') == 206
             and re.search(r'(?im)^content-range:\s*bytes\s+0-\d+/\d+', header) is not None
             and 0 < size <= limit
             and elapsed > 0 and math.isfinite(speed) and speed > 0)
    valid = response_ok and elapsed >= 1 and size >= config['min_bytes_per_sample']
    # A real 206 that transfers <1 MiB over the full window is slow evidence too.
    # A very fast sample hitting the byte cap in <1s also disproves a slow link.
    quality_valid = response_ok and (size >= config['min_bytes_per_sample'] or elapsed >= duration * .9)
    return {'provider': provider['name'], 'family': family, 'interface': iface,
            'remote_ip': data.get('remote_ip'), 'local_ip': data.get('local_ip'),
            'http_code': data.get('http_code'), 'bytes': size, 'seconds': elapsed,
            'mbps': round(speed * 8 / 1e6, 3), 'valid': valid,
            'quality_valid': quality_valid, 'curl_exit': p.returncode,
            'error': p.stderr.strip() if not valid else ''}


def main(recovery=False):
    try:
        run_main(recovery)
    except GameActivity as exc:
        log(str(exc) + '; previous weights and quarantine retained')


def run_main(recovery=False):
    STATE.parent.mkdir(parents=True, exist_ok=True)
    with open('/run/rpi-dualwan-speedtest.lock', 'w') as lock:
        try:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            log('another speed test is running; skipped')
            return
        config = json.loads(CONFIG.read_text())
        settings = policy(config)
        old = read_state(STATE)
        started = time.time()
        quarantine = old.get('quarantine', {})
        targets = {(family, iface) for family in (4, 6) for iface in IFACES
                   if not recovery or (quarantine.get(f'ipv{family}', {}).get(iface, {}).get('blocked') is True
                   and quarantine[f'ipv{family}'][iface].get('next_test_epoch', 0) <= started)}
        if not targets:
            log('no quarantined uplink due for recovery test; no download')
            return
        game_guard(config)
        if not recovery:
            interval = reserve_full_test(started)
            if not interval['allowed']:
                log(f"full speed test deferred: 2h minimum interval; {interval['remaining_seconds']}s remaining; no download")
                return
        sequence = int(old.get('sequence', 0)) + 1
        result = dict(old)
        report = []
        for family in (4, 6):
            ratios = {}
            selected = [iface for iface in IFACES if (family, iface) in targets]
            if not selected:
                continue
            for index, provider in enumerate(config['providers']):
                game_guard(config)
                if family not in provider.get('families', [4, 6]):
                    report.append({'family': family, 'provider': provider['name'],
                                   'skip_reason': 'family_disabled', 'skipped': 'address family not configured'})
                    continue
                available = [iface for iface in selected if healthy(family, iface)]
                for iface in selected:
                    if iface not in available:
                        report.append({'family': family, 'provider': provider['name'], 'interface': iface,
                                       'skip_reason': 'unhealthy', 'skipped': 'uplink not reachable'})
                if not available:
                    continue
                try:
                    host = urlsplit(provider['url']).hostname
                    candidates = addresses(host, family)
                    # Same destination IP and object for both interfaces in each pair.
                    address = candidates[0]
                    order = IFACES if (sequence + index) % 2 == 0 else IFACES[::-1]
                    pair = {}
                    for iface in order:
                        if iface not in available:
                            continue
                        try:
                            game_guard(config)
                            measurement = measure(provider, family, iface, address, config)
                        except (OSError, ValueError, subprocess.TimeoutExpired) as exc:
                            measurement = {'provider': provider['name'], 'family': family,
                                           'interface': iface, 'valid': False, 'error': str(exc)}
                        report.append(measurement)
                        pair[iface] = measurement
                        log(json.dumps(measurement))
                    if not recovery and all(pair.get(i, {}).get('valid') for i in IFACES) and healthy(family):
                        speeds = [pair[i]['mbps'] for i in IFACES]
                        if min(speeds) > 0:
                            ratios[provider['name']] = speeds[0] / sum(speeds)
                except (OSError, ValueError, IndexError) as exc:
                    for iface in available:
                        item = {'provider': provider['name'], 'family': family, 'interface': iface,
                                'skipped': str(exc), 'skip_reason': 'dns_unavailable'}
                        report.append(item)
                        log(json.dumps(item))
            if ratios:
                key = f'ipv{family}'
                result[key] = smooth(old.get(key, {}), list(ratios.values()),
                                     config['ewma_alpha'], config['deadband'], config['max_step'])
                result[key]['provider_ratios'] = ratios
                result[key]['last_success'] = time.strftime('%Y-%m-%dT%H:%M:%S%z')
                log('proposed ' + key + ' ' + json.dumps(result[key]))
            elif not recovery:
                log(f'IPv{family}: no valid pairs; retain previous weights')
        game_guard(config)
        finished = time.time()
        for item in report:
            item['tested_at'] = finished
        quarantine = {key: dict(value) for key, value in quarantine.items()}
        for family, iface in sorted(targets):
            states = quarantine.setdefault(f'ipv{family}', {})
            before = states.get(iface, {})
            evidence = [s for s in report if s.get('family') == family and s.get('interface') == iface]
            # Lost reachability during the run must not confirm a recovery.
            states[iface] = update_quality(before, evidence if healthy(family, iface) else [], settings, finished)
            log(f'quality ipv{family} {iface} ' + json.dumps(states[iface]))
        # Recovery probes replace only tested interface/family rows. Other rows retain their timestamps.
        samples = report
        if recovery:
            samples = [s for s in old.get('samples', [])
                       if ((s.get('interface') in IFACES and (s.get('family'), s['interface']) not in targets)
                           or (not s.get('interface') and s.get('family') not in {f for f, i in targets}))] + report
        result.update(sequence=sequence, updated=time.strftime('%Y-%m-%dT%H:%M:%S%z'),
                      elapsed_seconds=round(finished - started, 2), samples=samples,
                      mode='recovery' if recovery else 'full', quarantine=quarantine, quarantine_policy=settings,
                      downloaded_bytes=sum(item.get('bytes', 0) for item in report))
        with tempfile.NamedTemporaryFile(mode='w', dir=STATE.parent, prefix='.weights-', delete=False) as out:
            json.dump(result, out, indent=2)
            out.write('\n')
            out.flush()
            os.fsync(out.fileno())
            temporary = pathlib.Path(out.name)
        temporary.chmod(0o644)
        temporary.replace(STATE)
        if not recovery:
            archive(report, finished)
        log('saved weights; controller will apply within the next health cycle')


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--recovery', action='store_true', help='test only quarantined uplinks due for a retest')
    parser.add_argument('--check-game-guard', action='store_true', help='check game activity without downloading or changing weights')
    options = parser.parse_args()
    if options.check_game_guard:
        try:
            game_guard(json.loads(CONFIG.read_text()))
            log('game guard clear; no download performed')
        except GameActivity as exc:
            log(str(exc))
    else:
        main(recovery=options.recovery)
