from argparse import ArgumentParser
import ipaddress
import os
from os.path import join
import sys
import geoip2.database
import netaddr


def try_parse(ip):
    try:
        return ipaddress.ip_network(ip)
    except ValueError:
        return None


def get_home() -> str:
    if sys.platform == "win32":
        return join(os.environ["USERPROFILE"], "AppData", "Roaming")
    else:
        return os.environ["HOME"]


p = ArgumentParser()
p.add_argument("CONFIG", help="Configuration files", nargs="+")
p.add_argument("-c", "--country", help="Country code", default="CN")
p.add_argument("-o", "--output", help="Output file",
               default="ip-cidr-detect.yaml")
p.add_argument("--resolve", help="Resolve domain name", action="store_true",
               default=False)
p.add_argument("-r", "--rule", help="Clash rule target", default="PROXY")
p.add_argument("-d", "--db", help="GeoIP database",
               default=join(get_home(), ".config/clash/Country.mmdb"))
arg = p.parse_intermixed_args()
with geoip2.database.Reader(arg.db) as reader:
    ips = []
    for config in arg.CONFIG:
        with open(config, "r", encoding="UTF-8") as i:
            tmp = i.readline()
            while tmp:
                tmp = tmp.split("#")[0].strip()
                print(tmp)
                sl = tmp.split("=")
                if len(sl) == 1 and not tmp.startswith("["):
                    network = try_parse(tmp)
                    if network is not None:
                        for ip in network:
                            re = reader.country(ip)
                            if re.country.iso_code == arg.country:
                                ips.append(str(ip))
                tmp = i.readline()
    ips = netaddr.cidr_merge(ips)
    print(ips)
    with open(arg.output, "w", encoding="UTF-8") as o:
        o.write("# Generated by ip-cidr-detect.py\n")
        o.write("rules:\n")
        for r in ips:
            o.write(f"- IP-CIDR,{r},{arg.rule}{'' if arg.resolve else ',no-resolve'}\n")  # noqa: E501