#!/usr/bin/env python3
"""HY2 Panel AnyTLS agent. Local-only core APIs; HTTPS outbound control plane."""
import fcntl
import json
import os
from pathlib import Path
import subprocess
import sys
import time
import urllib.request
import urllib.error
import grpc
import stats_pb2

ROOT = Path('/etc/hy2-panel-anytls')
STATE = Path('/var/lib/hy2-panel-anytls')
BINARY = '/usr/local/lib/hy2-panel-anytls/sing-box'
UNIT = 'hy2-panel-anytls.service'


def atomic_json(file, value):
    temporary = file.with_suffix('.tmp')
    with open(temporary, 'w') as handle:
        os.chmod(temporary, 0o600)
        json.dump(value, handle, sort_keys=True)
        handle.flush()
        os.fsync(handle.fileno())
    os.replace(temporary, file)


def core(command):
    subprocess.run(['systemctl', command, UNIT], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def read_json(file):
    return json.loads(file.read_text())


def epoch():
    pid = subprocess.check_output(['systemctl', 'show', UNIT, '-p', 'MainPID', '--value'], text=True).strip()
    if pid == '0':
        return None
    # PID + boot ID + start ticks survives agent restarts without re-counting.
    start = Path('/proc/' + pid + '/stat').read_text().rsplit(')', 1)[1].split()[19]
    return Path('/proc/sys/kernel/random/boot_id').read_text().strip() + ':' + pid + ':' + start


def snapshot():
    current = epoch()
    if current is None:
        return None
    config = read_json(ROOT / 'config.json')
    experimental = config['experimental']
    with grpc.insecure_channel(experimental['v2ray_api']['listen']) as channel:
        query = channel.unary_unary(
            '/v2ray.core.app.stats.command.StatsService/QueryStats',
            request_serializer=stats_pb2.QueryStatsRequest.SerializeToString,
            response_deserializer=stats_pb2.QueryStatsResponse.FromString)
        result = query(stats_pb2.QueryStatsRequest(pattern='user>>>', reset=False), timeout=5)
    traffic = {}
    for stat in result.stat:
        parts = stat.name.split('>>>')
        if len(parts) == 4 and parts[0] == 'user' and parts[2] == 'traffic':
            value = traffic.setdefault(parts[1], {'tx': 0, 'rx': 0})
            # HY2 ledgers use tx=server->client, rx=client->server.
            if parts[3] == 'uplink':
                value['rx'] = int(stat.value)
            elif parts[3] == 'downlink':
                value['tx'] = int(stat.value)
    clash = experimental['clash_api']
    request = urllib.request.Request('http://' + clash['external_controller'] + '/connections',
                                     headers={'Authorization': 'Bearer ' + clash['secret']})
    with urllib.request.urlopen(request, timeout=5) as response:
        streams = len(json.load(response).get('connections', []))
    if epoch() != current:
        raise RuntimeError('core_restarted_during_snapshot')
    return {'epoch': current, 'traffic': traffic, 'streamCount': streams}


def remember(state):
    sample = snapshot()
    if sample:
        state['pending'][sample['epoch']] = sample
        atomic_json(STATE / 'state.json', state)
    return sample


def sync(credentials, sample):
    request = urllib.request.Request(credentials['panel'] + '/api/agent/' + credentials['slug'] + '/sync',
        data=json.dumps(sample).encode(), method='POST',
        headers={'Content-Type': 'application/json', 'Authorization': 'Bearer ' + credentials['token']})
    with urllib.request.urlopen(request, timeout=10) as response:
        return json.load(response)


def apply_config(config, state):
    previous = read_json(ROOT / 'config.json')
    if previous != config:
        # Capture one last cumulative sample before replacing a core process.
        remember(state)
        staged = ROOT / 'next.json'
        atomic_json(staged, config)
        checked = subprocess.run([BINARY, 'check', '-c', str(staged)], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        if checked.returncode:
            core('stop')
            raise RuntimeError('invalid_core_configuration')
        os.replace(staged, ROOT / 'config.json')
        core('restart')
    else:
        core('start')


def run():
    STATE.mkdir(mode=0o700, parents=True, exist_ok=True)
    lock = open(STATE / 'agent.lock', 'w')
    fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
    credentials = read_json(ROOT / 'agent.json')
    state_file = STATE / 'state.json'
    state = read_json(state_file) if state_file.exists() else {'pending': {}}
    # No cached lease is trusted after an agent restart.
    deadline = 0
    while True:
        try:
            if time.monotonic() >= deadline:
                try:
                    remember(state)
                finally:
                    core('stop')
            else:
                remember(state)
            if len(state['pending']) > 256:
                core('stop')
                raise RuntimeError('report_backlog_full')
            reports = list(state['pending'].items())
            if not reports:
                reports = [('heartbeat', {'epoch': 'heartbeat', 'traffic': {}, 'streamCount': 0})]
            desired = None
            for report_epoch, report in reports:
                desired = sync(credentials, report)
                state['pending'].pop(report_epoch, None)
                atomic_json(state_file, state)
            apply_config(desired['config'], state)
            deadline = time.monotonic() + min(90, max(1, desired['leaseSeconds']))
            print('同步成功', flush=True)
        except Exception as error:
            # Do not log credentials, URLs, user lists or config contents.
            print('同步失败：' + type(error).__name__, flush=True)
            if isinstance(error, urllib.error.HTTPError) and error.code in (400, 401, 403, 404):
                deadline = 0
                core('stop')
            if time.monotonic() >= deadline:
                core('stop')
        time.sleep(max(1, min(15, deadline - time.monotonic())) if deadline else 5)


if __name__ == '__main__':
    os.umask(0o077)
    try:
        run()
    except Exception as error:
        core('stop')
        print('管理程序停止：' + type(error).__name__, flush=True)
        sys.exit(1)