#!/usr/bin/env python
"""
@Author: xiaobaiTser
@Time  : 2023/8/15 0:51
@File  : FlushDNSUtils.py
@Desc  : GitHub DNS 解析工具 - 多线程获取并更新 hosts 文件
"""
import os
import sys
import time
import json
import ctypes
import subprocess
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime

try:
    import dns.resolver
except ImportError:
    print("正在安装 dnspython...")
    subprocess.run([sys.executable, "-m", "pip", "install", "dnspython", "-q"], check=True)
    import dns.resolver

try:
    import requests
except ImportError:
    print("正在安装 requests...")
    subprocess.run([sys.executable, "-m", "pip", "install", "requests", "-q"], check=True)
    import requests

try:
    from tkinter.messagebox import showerror, showinfo
except ModuleNotFoundError:
    showinfo = showerror = print

GITHUB_DOMAINS = [
    "github.com",
    "raw.githubusercontent.com",
    "assets-cdn.github.com",
    "github.global.ssl.fastly.net",
    "codeload.github.com",
    "objects.githubusercontent.com",
]

PLATFORM_HOSTS = {
    "win32": r"C:\Windows\System32\drivers\etc\hosts",
    "linux": "/etc/hosts",
    "darwin": "/etc/hosts",
}


def get_dns_ips(domain: str, timeout: float = 5) -> list[str]:
    """获取域名的 A 记录 IP"""
    ips = []
    try:
        resolver = dns.resolver.Resolver()
        resolver.timeout = timeout
        resolver.lifetime = timeout
        answers = resolver.resolve(domain, "A")
        for answer in answers:
            ip = answer.to_text()
            if ip != "0.0.0.0":
                ips.append(ip)
    except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer, dns.resolver.Timeout) as e:
        print(f"  [警告] {domain}: DNS 查询失败 - {e}")
    return ips


def get_myssl_ip(domain: str, timeout: float = 5) -> list[str]:
    """通过 MySSL API 获取域名 IP"""
    headers = {
        "Accept": "application/json, text/javascript, */*; q=0.01",
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
        "Referer": "https://myssl.com/dns_check.html",
    }
    params = {"qtype": "1", "host": domain, "qmode": "-1"}
    try:
        resp = requests.get(
            "https://myssl.com/api/v1/tools/dns_query",
            params=params,
            headers=headers,
            timeout=timeout,
        ).json()
        ips = []
        for record_type in ["01", "852", "86"]:
            try:
                for ip in resp["data"][record_type]["answer"]["records"]:
                    if ip["value"] != "0.0.0.0":
                        ips.append(ip["value"])
            except (KeyError, TypeError):
                continue
        return list(set(ips))
    except (requests.RequestException, json.JSONDecodeError) as e:
        print(f"  [警告] {domain}: MySSL API 失败 - {e}")
        return []


def get_ip138_ip(domain: str, timeout: float = 5) -> list[str]:
    """通过 IP138 API 获取域名 IP"""
    headers = {
        "Accept": "*/*",
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
        "Referer": f"https://site.ip138.com/{domain}/",
    }
    params = {"domain": domain, "time": int(time.time() * 1000)}
    try:
        resp = requests.get(
            "https://site.ip138.com/domain/read.do",
            params=params,
            headers=headers,
            timeout=timeout,
        ).json()
        if resp.get("status"):
            return [ip["ip"] for ip in resp["data"] if ip["ip"] != "0.0.0.0"]
    except (requests.RequestException, json.JSONDecodeError) as e:
        print(f"  [警告] {domain}: IP138 API 失败 - {e}")
    return []


def resolve_domain(domain: str) -> dict[str, list[str]]:
    """多策略解析域名 IP"""
    results = {}

    ips = get_dns_ips(domain)
    if ips:
        results["dns"] = ips

    if not ips:
        ips = get_myssl_ip(domain)
        if ips:
            results["myssl"] = ips

    if not ips:
        ips = get_ip138_ip(domain)
        if ips:
            results["ip138"] = ips

    return {domain: results.get("dns", results.get("myssl", results.get("ip138", [])))}


def ping_ip(ip: str, count: int = 2, timeout: int = 1000) -> int:
    """测试 IP 延迟，返回毫秒数"""
    try:
        if sys.platform == "win32":
            output = subprocess.run(
                ["ping", "-n", str(count), "-w", str(timeout), ip],
                capture_output=True,
                text=True,
                timeout=10,
            )
            for line in output.stdout.split("\n"):
                if "平均 =" in line or "Average =" in line:
                    import re
                    match = re.search(r"平均 = (\d+)ms|Average = (\d+)ms", line)
                    if match:
                        return int(match.group(1)) if match.group(1) else int(match.group(2))
                if "时间 =" in line or "time =" in line:
                    import re
                    match = re.search(r"时间=(\d+)ms|time=(\d+)ms", line)
                    if match:
                        return int(match.group(1)) if match.group(1) else int(match.group(2))
        else:
            output = subprocess.run(
                ["ping", "-c", str(count), "-W", str(timeout // 1000), ip],
                capture_output=True,
                text=True,
                timeout=10,
            )
            import re
            match = re.search(r"min/avg/max.*= [\d.]+/([\d.]+)/", output.stdout)
            if match:
                return int(float(match.group(1)))
    except Exception:
        pass
    return timeout


def test_ip_latency(dns_results: dict[str, list[str]], max_ips: int = 2) -> dict[str, list[str]]:
    """测试 IP 延迟，返回每个域名最快的 IP（最多 max_ips 个）"""
    print(f"\n{'='*60}")
    print(f"  测试 IP 延迟 (选择最快的 {max_ips} 个 IP)")
    print(f"{'='*60}\n")

    for domain, ips in dns_results.items():
        if not ips:
            continue

        print(f"  {domain}:")
        ip_latencies = []
        for ip in ips:
            latency = ping_ip(ip)
            ip_latencies.append((ip, latency))
            status = "✓" if latency < 200 else "△" if latency < 500 else "✗"
            print(f"    {status} {ip}: {latency}ms")

        ip_latencies.sort(key=lambda x: x[1])
        best_ips = [ip for ip, _ in ip_latencies[:max_ips]]
        dns_results[domain] = best_ips
        print(f"    → 保留: {', '.join(best_ips)}")

    return dns_results


def resolve_domains_multi_thread(domains: list[str], max_workers: int = 5) -> dict[str, list[str]]:
    """多线程解析域名"""
    all_results = {}
    total = len(domains)

    print(f"\n{'='*60}")
    print(f"  开始解析 {total} 个域名 (使用 {max_workers} 个线程)")
    print(f"{'='*60}\n")

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_domain = {executor.submit(resolve_domain, d): d for d in domains}

        completed = 0
        for future in as_completed(future_to_domain):
            domain = future_to_domain[future]
            completed += 1
            progress = f"[{completed}/{total}] {progress_bar(completed, total)} {completed * 100 // total}%"

            try:
                result = future.result()
                if result[domain]:
                    all_results.update(result)
                    print(f"✓ {progress} {domain}: {', '.join(result[domain])}")
                else:
                    print(f"✗ {progress} {domain}: 未能解析到 IP")
            except Exception as e:
                print(f"✗ {progress} {domain}: 解析异常 - {e}")

    return all_results


def progress_bar(current: int, total: int, length: int = 30) -> str:
    """生成进度条"""
    percent = current / total
    filled = int(length * percent)
    bar = "█" * filled + "░" * (length - filled)
    return bar


def read_hosts_file() -> list[str]:
    """读取 hosts 文件内容"""
    hosts_path = PLATFORM_HOSTS.get(sys.platform, PLATFORM_HOSTS["win32"])

    if not os.path.exists(hosts_path):
        print(f"[错误] hosts 文件不存在: {hosts_path}")
        return []

    try:
        with open(hosts_path, "r", encoding="utf-8") as f:
            lines = f.readlines()
    except PermissionError:
        print(f"[错误] 无权限读取: {hosts_path}")
        return []

    return [line.strip() for line in lines]


def update_hosts(dns_results: dict[str, list[str]], domains: list[str]) -> bool:
    """更新 hosts 文件 - 只修改指定域名，保留其他所有内容"""
    hosts_path = PLATFORM_HOSTS.get(sys.platform, PLATFORM_HOSTS["win32"])
    original_lines = read_hosts_file()

    domains_to_update = set(domains)

    new_content = []

    for line in original_lines:
        stripped = line.strip()
        
        if not stripped or stripped.startswith("#"):
            new_content.append(stripped)
            continue
        
        parts = stripped.split()
        if len(parts) >= 2:
            domain_part = " ".join(parts[1:])
            if domain_part in domains_to_update:
                if domain_part in dns_results and dns_results[domain_part]:
                    for ip in dns_results[domain_part]:
                        entry = f"{ip}\t{domain_part}"
                        new_content.append(entry)
                        print(f"  → 更新: {entry}")
                continue
        
        new_content.append(stripped)

    for domain in domains_to_update:
        if domain not in [line.split()[1] if len(line.split()) >= 2 else "" for line in new_content if line.strip() and not line.strip().startswith("#")]:
            if domain in dns_results and dns_results[domain]:
                for ip in dns_results[domain]:
                    entry = f"{ip}\t{domain}"
                    new_content.append(entry)
                    print(f"  → 追加: {entry}")

    try:
        with open(hosts_path, "w", encoding="utf-8") as f:
            f.write("\n".join(new_content))
        print(f"\n✓ hosts 文件已更新: {hosts_path}")
        return True
    except PermissionError:
        print(f"\n[错误] 无权限写入: {hosts_path}")
        print("  提示: 以管理员身份运行此脚本")
        return False


def flush_dns_cache() -> bool:
    """刷新 DNS 缓存"""
    print(f"\n{'='*60}")
    print("  刷新 DNS 缓存")
    print(f"{'='*60}\n")

    if sys.platform == "win32":
        try:
            os.system("ipconfig /flushdns >nul 2>&1")
            print("✓ Windows DNS 缓存已刷新")
            return True
        except Exception as e:
            print(f"✗ 刷新失败: {e}")
            return False
    elif sys.platform == "darwin":
        try:
            os.system("sudo dscacheutil -flushcache > /dev/null 2>&1")
            os.system("sudo killall -HUP mDNSResponder > /dev/null 2>&1")
            print("✓ macOS DNS 缓存已刷新")
            return True
        except Exception as e:
            print(f"✗ 刷新失败: {e}")
            return False
    elif sys.platform == "linux":
        try:
            os.system("sudo systemd-resolve --flush-caches > /dev/null 2>&1")
            os.system("sudo resolvectl flush-caches > /dev/null 2>&1")
            print("✓ Linux DNS 缓存已刷新")
            return True
        except Exception as e:
            print(f"✗ 刷新失败: {e}")
            return False
    else:
        print(f"⚠ 不支持的平台: {sys.platform}")
        return False


def flush_dns(domains: list[str] = None, show_ui: bool = True) -> bool:
    """主函数：解析 DNS 并更新 hosts"""
    if domains is None:
        domains = GITHUB_DOMAINS.copy()

    print(f"\n{'='*60}")
    print(f"  GitHub DNS 解析工具")
    print(f"  时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"{'='*60}")

    dns_results = resolve_domains_multi_thread(domains)

    if not any(dns_results.values()):
        print("\n[错误] 所有域名均未能解析，请检查网络连接")
        if show_ui:
            showerror("错误", "所有域名均未能解析，请检查网络连接")
        return False

    dns_results = test_ip_latency(dns_results, max_ips=2)

    if not update_hosts(dns_results, domains.copy()):
        return False

    if not flush_dns_cache():
        print("⚠ DNS 缓存刷新失败，但 hosts 已更新")

    print(f"\n{'='*60}")
    print("  完成!")
    print(f"{'='*60}\n")

    if show_ui:
        showinfo("完成", "hosts 文件已更新，DNS 缓存已刷新")

    return True


def run_as_admin() -> None:
    """请求管理员权限 (Windows)"""
    if sys.platform == "win32" and not ctypes.windll.shell32.IsUserAnAdmin():
        print("正在请求管理员权限...")
        ctypes.windll.shell32.ShellExecuteW(
            None, "runas", sys.executable, " ".join(sys.argv), None, 1
        )
        sys.exit()

def main() -> None:
    """主函数
    刷新 GitHub 域名的 DNS 记录并更新 hosts 文件。
    
    示例:
        xiaobaicmd tools dns github.com,raw.githubusercontent.com
    """
    if len(sys.argv) > 1 and sys.argv[1] == "dns":
        domains = sys.argv[2:] if len(sys.argv) > 2 else GITHUB_DOMAINS.copy()
        flush_dns(domains, show_ui=False)
    else:
        run_as_admin()
        flush_dns(show_ui=True)

