#!/usr/bin/env python3
import argparse
import subprocess
import time
import sys
import os
import shutil
import tempfile
import re
from pathlib import Path

def parse_args():
    p = argparse.ArgumentParser(
        description="Проверка VLESS-конфигов через sing-box + curl"
    )
    p.add_argument("indir", type=Path,
                   help="Папка с JSON-конфигами (VLESS)")
    p.add_argument("--port", "-p", type=int, default=5055,
                   help="Локальный порт для sing-box (по умолчанию 5055)")
    p.add_argument("--test-url", "-u", default="http://example.com",
                   help="URL для проверки (curl --proxy)")
    p.add_argument("--trusted", "-t", type=Path, default=Path("./trusted"),
                   help="Куда копировать успешные конфиги (по умолчанию ./trusted)")
    p.add_argument("--copy-ok", action="store_true",
                   help="Копировать успешные конфиги в папку trusted")
    p.add_argument("--timeout", "-T", type=int, default=10,
                   help="Таймаут на стартап sing-box и на curl (в секундах)")
    p.add_argument("--latency", action="store_true",
                   help="Измерять TCP connect latency (ms)")
    p.add_argument("--bandwidth", action="store_true",
                   help="Измерять скорость скачивания (KB/s)")
    return p.parse_args()


def prepare_config_with_port(orig_path: Path, port: int) -> Path:
    text = orig_path.read_text(encoding="utf-8")
    new_text, count = re.subn(
        r'"listen_port"\s*:\s*\d+',
        f'"listen_port": {port}',
        text,
        count=1
    )
    if count == 0:
        new_text = text

    tmp = tempfile.NamedTemporaryFile(
        mode="w", delete=False, suffix=".json", encoding="utf-8"
    )
    tmp.write(new_text)
    tmp.flush()
    tmp.close()
    return Path(tmp.name)


def check_config(orig_config: Path, port: int, test_url: str,
                 timeout: int, do_latency: bool, do_bw: bool):
    """
    Возвращает: (ok: bool, latency_ms: float|None, bw_kb: float|None)
    """
    temp_config = prepare_config_with_port(orig_config, port)

    cmd = ["./sing-box", "run", "-c", str(temp_config)]
    try:
        proc = subprocess.Popen(cmd,
                                stdout=subprocess.DEVNULL,
                                stderr=subprocess.DEVNULL)
    except FileNotFoundError:
        print("ERROR: команда 'sing-box' не найдена в PATH", file=sys.stderr)
        temp_config.unlink(missing_ok=True)
        return False, None, None

    time.sleep(1)

    # Формируем curl
    curl_cmd = [
        "curl",
        "--proxy", f"socks5h://127.0.0.1:{port}",
        "--max-time", str(timeout),
        "--silent",
        "--output", os.devnull,
    ]

    write_parts = []
    if do_latency:
        write_parts.append("%{time_connect}")
    if do_bw:
        write_parts.append("%{speed_download}")
    # всегда нужен write-out, чтобы curl не выдавал дифолтный progress
    if write_parts:
        curl_cmd += ["--write-out", " ".join(write_parts)]
    # и сам URL
    curl_cmd.append(test_url)

    try:
        result = subprocess.run(
            curl_cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            timeout=timeout,
            universal_newlines=True
        )
        ok = (result.returncode == 0)
        out = result.stdout.strip()
        latency_ms = None
        bw_kb = None
        if ok and out and write_parts:
            parts = out.split()
            idx = 0
            if do_latency:
                # time_connect в секундах
                try:
                    t = float(parts[idx])
                    latency_ms = t * 1000.0
                except:
                    latency_ms = None
                idx += 1
            if do_bw:
                # speed_download в байтах/сек -> KB/s
                try:
                    sp = float(parts[idx])
                    bw_kb = sp / 1024.0
                except:
                    bw_kb = None
    except subprocess.TimeoutExpired:
        ok = False
        latency_ms = None
        bw_kb = None
    finally:
        proc.terminate()
        try:
            proc.wait(2)
        except subprocess.TimeoutExpired:
            proc.kill()
        temp_config.unlink(missing_ok=True)

    return ok, latency_ms, bw_kb


def main():
    args = parse_args()

    if not args.indir.exists() or not args.indir.is_dir():
        print(f"ERROR: папка {args.indir} не найдена или не каталог", file=sys.stderr)
        sys.exit(1)

    if args.copy_ok:
        args.trusted.mkdir(parents=True, exist_ok=True)

    for cfg in sorted(args.indir.glob("*.json")):
        ok, latency_ms, bw_kb = check_config(
            cfg,
            port=args.port,
            test_url=args.test_url,
            timeout=args.timeout,
            do_latency=args.latency,
            do_bw=args.bandwidth
        )

        parts = []
        parts.append("OK" if ok else "FAILED")
        if ok:
            if args.latency and latency_ms is not None:
                parts.append(f"latency={latency_ms:.1f}ms")
            if args.bandwidth and bw_kb is not None:
                parts.append(f"bw={bw_kb:.1f}KB/s")

        print(f"{cfg.name}: " + " ".join(parts))

        if ok and args.copy_ok:
            shutil.copy2(cfg, args.trusted / cfg.name)


if __name__ == "__main__":
    main()