#!/usr/bin/python3
"""Reconcile MPTCP quarantine using validated, locked Mihomo hot reloads."""
import copy
import fcntl
import json
import os
import pathlib
import subprocess
import sys
import tempfile
import urllib.error
import urllib.parse
import urllib.request
import yaml
sys.path.insert(0, '/usr/local/lib')
from mihomo_local import apply_mptcp_nodes
from mihomo_mptcp_policy import load_runtime


def atomic(path, content):
    fd, name = tempfile.mkstemp(prefix='.' + path.name, dir=path.parent)
    try:
        with os.fdopen(fd, 'w') as out:
            out.write(content)
        os.replace(name, path)
    finally:
        pathlib.Path(name).unlink(missing_ok=True)


def reconcile(base=pathlib.Path('/etc/mihomo'), policy=pathlib.Path('/etc/mihomo/local-routing.json'),
              state=pathlib.Path('/var/lib/mihomo-update/mptcp-guard.json'), runtime=None,
              controller='http://127.0.0.1:9090', binary='/usr/local/bin/mihomo'):
    path = base / 'config.yaml'
    pending = base / '.mptcp-guard-pending'
    original = path.read_text()
    cfg = yaml.load(original, Loader=getattr(yaml, 'CSafeLoader', yaml.SafeLoader))
    names = json.loads(policy.read_text()).get('mptcp_nodes', [])
    candidate = copy.deepcopy(cfg)
    runtime = load_runtime() if runtime is None else runtime
    apply_mptcp_nodes(candidate, names, runtime)
    old = {p['name']: p for p in cfg.get('proxies', []) if p.get('name') in names}
    new = {p['name']: p for p in candidate.get('proxies', []) if p.get('name') in names}
    changed = [n for n in names if old[n] != new[n]]
    modes = {n: 'mptcp' if p['mptcp'] else 'tcp-' + p['interface-name'] for n, p in new.items()}
    report = {'modes': modes, 'changed': changed, 'status': 'unchanged'}
    if not changed and not pending.exists():
        # Do not rewrite the state or config on each timer tick.
        prior = json.loads(state.read_text()) if state.exists() else {}
        if prior.get('modes') != modes:
            atomic(state, json.dumps(report, ensure_ascii=False, indent=2) + '\n')
        return report
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))

    def api(method, endpoint, body=None):
        data = json.dumps(body).encode() if body is not None else None
        request = urllib.request.Request(controller + endpoint, data=data, method=method,
            headers={'Authorization': 'Bearer ' + cfg.get('secret', ''), 'Content-Type': 'application/json'})
        with opener.open(request, timeout=15) as response:
            content = response.read()
            return json.loads(content) if content else {}

    selectors = {n: p['now'] for n, p in api('GET', '/proxies').get('proxies', {}).items()
                 if p.get('type') == 'Selector' and p.get('now')}

    def reload_config():
        api('PUT', '/configs?force=false', {'path': str(path)})
        for name, choice in selectors.items():
            api('PUT', '/proxies/' + urllib.parse.quote(name, safe=''), {'name': choice})

    # Existing connections retain their old adapter and socket; never DELETE
    # connections, change kernel MPTCP state, or remove endpoints here.
    text = yaml.dump(candidate, Dumper=getattr(yaml, 'CSafeDumper', yaml.SafeDumper),
                     allow_unicode=True, sort_keys=False)
    with tempfile.TemporaryDirectory(prefix='.mptcp-guard-', dir=base) as temp:
        staged = pathlib.Path(temp) / 'config.yaml'
        staged.write_text(text)
        result = subprocess.run([binary, '-t', '-d', str(base), '-f', str(staged)],
                                capture_output=True, timeout=45,
                                env={**os.environ, 'SAFE_PATHS': str(base)})
        if result.returncode:
            raise RuntimeError('Mihomo rejected MPTCP policy candidate')
    # Save recovery intent first, so interruption after the disk write is retried.
    atomic(pending, json.dumps({'modes': modes}))
    atomic(path, text)
    try:
        reload_config()
        pending.unlink()
    except Exception:
        # Restore disk and live state on failure; pending intent survives for retry.
        atomic(path, original)
        try:
            reload_config()
        except Exception:
            pass
        raise
    import datetime
    report.update(status='applied', existing_connections='preserved',
                  updated=datetime.datetime.now().astimezone().isoformat())
    atomic(state, json.dumps(report, ensure_ascii=False, indent=2) + '\n')
    print(json.dumps(report, ensure_ascii=False), flush=True)
    return report


def main():
    os.umask(0o077)
    with open('/run/mihomo-update.lock', 'a') as lock:
        try:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            return  # retry after the subscription/core update has completed
        reconcile()


if __name__ == '__main__':
    try:
        main()
    except Exception as error:
        # Never print YAML, controller credentials or urllib request bodies.
        print('MPTCP quarantine reconciliation failed: ' + type(error).__name__, file=sys.stderr)
        sys.exit(1)
