#!/usr/bin/env python3

import sys
import urllib.request
import os

categories = {
    "ads": [
        ("EasyList", "https://easylist.to/easylist/easylist.txt"),
        (
            "uBlock",
            "https://raw.githubusercontent.com/uBlockOrigin/uAssets/master/filters/filters.txt",
        ),
        (
            "AdGuard Base",
            "https://raw.githubusercontent.com/AdguardTeam/FiltersRegistry/master/filters/filter_2_Base/filter.txt",
        ),
        ("Blocklist Project Ads", "https://blocklistproject.github.io/Lists/ads.txt"),
        (
            "GoodbyeAds",
            "https://raw.githubusercontent.com/jerryn70/GoodbyeAds/master/Formats/GoodbyeAds-AdBlock-Filter.txt",
        ),
    ],
    "trackers": [
        ("EasyPrivacy", "https://easylist.to/easylist/easyprivacy.txt"),
        (
            "AdGuard DNS",
            "https://adguardteam.github.io/AdGuardSDNSFilter/Filters/filter.txt",
        ),
        (
            "Blocklist Project Tracking",
            "https://blocklistproject.github.io/Lists/tracking.txt",
        ),
        ("OISD big", "https://raw.githubusercontent.com/sjhgvr/oisd/main/oisd_big.txt"),
    ],
    "porn": [
        (
            "StevenBlack porn",
            "https://raw.githubusercontent.com/StevenBlack/hosts/master/alternates/porn-only/hosts",
        ),
        ("Blocklist Project Porn", "https://blocklistproject.github.io/Lists/porn.txt"),
        (
            "OISD nsfw",
            "https://raw.githubusercontent.com/sjhgvr/oisd/main/abp_nsfw.txt",
        ),
    ],
    "gambling": [
        (
            "StevenBlack gambling",
            "https://raw.githubusercontent.com/StevenBlack/hosts/master/alternates/gambling-only/hosts",
        ),
        (
            "HaGeZi Gambling",
            "https://raw.githubusercontent.com/hagezi/dns-blocklists/main/adblock/gambling.txt",
        ),
        (
            "Blocklist Project Gambling",
            "https://blocklistproject.github.io/Lists/gambling.txt",
        ),
    ],
    "malware": [
        (
            "HaGeZi TIF",
            "https://raw.githubusercontent.com/hagezi/dns-blocklists/main/adblock/tif.txt",
        ),
        (
            "Blocklist Project Malware",
            "https://blocklistproject.github.io/Lists/malware.txt",
        ),
        (
            "Blocklist Project Phishing",
            "https://blocklistproject.github.io/Lists/phishing.txt",
        ),
    ],
}

BLOCKLIST_DIR = "/etc/linum/blocklist"
CUSTOM_FILE = os.path.join(BLOCKLIST_DIR, "custom.txt")
WHITELIST_FILE = os.path.join(BLOCKLIST_DIR, "whitelist.txt")


def eprint(*a, **kw):
    print(*a, file=sys.stderr, **kw)


def confirm(prompt):
    return input(f"{prompt} [Y/n] ").strip().lower() in ("", "y", "yes")


def download_text(url):
    try:
        with urllib.request.urlopen(url, timeout=30) as resp:
            text = resp.read().decode("utf-8", errors="replace")
        return text.splitlines()
    except Exception as e:
        eprint(f"download failed {url}: {e}")
        return None


def cmd_list(args):
    if not args:
        for k in categories:
            print(k)
        return

    key = args[0]
    if key == "all":
        for k, items in categories.items():
            print(f"\n[{k}]")
            for name, url in items:
                print(f"-> {name}")
                print(f"-> {url}")
        return
    items = categories.get(key)
    if not items:
        eprint(f"'{key}' category not found")
        return
    for name, url in items:
        print(f">> {name}")
        print(f">> {url}")


def cmd_download(_args):
    cat_keys = list(categories)
    print("available:")
    for i, k in enumerate(cat_keys, 1):
        print(f">> {i}, {k}")
    raw = input("choose one category: ")
    sel_cats = []
    for part in raw.split():
        try:
            idx = int(part) - 1
            if 0 <= idx < len(cat_keys):
                sel_cats.append(cat_keys[idx])
        except ValueError:
            eprint(f"skip '{part}' because it is not a number")
    if not sel_cats:
        eprint("there is no valid category")
        return

    selected = []
    for cat in sel_cats:
        items = categories[cat]
        print(f"\nitem in [{cat}]:")
        for i, (name, url) in enumerate(items, 1):
            print(f">> {i}. {name}")
        raw2 = input(f"choose one from '{cat}' (space or '*' for all): ")
        if raw2.strip() == "*":
            selected.extend((name, url, f"{cat}.txt") for name, url in items)
        else:
            for part in raw2.split():
                try:
                    idx = int(part) - 1
                    if 0 <= idx < len(items):
                        name, url = items[idx]
                        selected.append((name, url, f"{cat}.txt"))
                except ValueError:
                    pass

    if not selected:
        eprint("nothing selected")
        return

    print(f"\nit will download {len(selected)} item and merge to the blocklist")
    if not confirm("continue?"):
        print("cancel")
        return

    file_buffers: dict[str, list[str]] = {}
    for name, url, fname in selected:
        eprint(f"download {name}...")
        lines = download_text(url)
        if lines is None:
            continue
        fpath = os.path.join(BLOCKLIST_DIR, fname)
        existing = []
        if os.path.exists(fpath):
            with open(fpath, "r") as f:
                existing = f.read().splitlines()
        all_lines = existing + lines
        seen = set()
        unique = []
        for ln in all_lines:
            stripped = ln.strip()
            if stripped and stripped not in seen:
                seen.add(stripped)
                unique.append(stripped)
        file_buffers[fname] = unique

    for fname, lines in file_buffers.items():
        fpath = os.path.join(BLOCKLIST_DIR, fname)
        old_count = len(open(fpath).read().splitlines()) if os.path.exists(fpath) else 0
        new_count = len(lines)
        delta = new_count - old_count
        eprint(f"  {fname}: {old_count} to {new_count} lines ({delta:+d})")

    if not confirm("write the file?"):
        print("cancel")
        return

    for fname, lines in file_buffers.items():
        fpath = os.path.join(BLOCKLIST_DIR, fname)
        with open(fpath, "w") as f:
            f.write("\n".join(lines) + "\n")
        print(f"{fname} — {len(lines)} lines")


def cmd_custom(args):
    if not os.path.exists(CUSTOM_FILE):
        open(CUSTOM_FILE, "w").close()

    with open(CUSTOM_FILE, "r") as f:
        lines = f.read().splitlines()

    if not args or args[0] == "list":
        if not lines:
            print("(custom.txt is empty)")
        for ln in lines:
            print(ln)
        return

    cmd = args[0]
    if cmd == "add" and len(args) >= 3:
        domain, ip = args[1], args[2]
        lines.append(f"{ip} {domain}")
    elif cmd == "del" and len(args) >= 2:
        domain = args[1]
        lines = [ln for ln in lines if not ln.endswith(f" {domain}")]
    elif cmd == "set" and len(args) >= 3:
        domain, new_ip = args[1], args[2]
        for i, ln in enumerate(lines):
            parts = ln.split()
            if len(parts) == 2 and parts[1] == domain:
                lines[i] = f"{new_ip} {domain}"
                break
        else:
            eprint(f"domain '{domain}' not found in custom.txt")
            return
    else:
        eprint("Usage: custom list|add <domain> <ip>|del <domain>|set <domain> <ip>")
        return

    with open(CUSTOM_FILE, "w") as f:
        f.write("\n".join(lines) + "\n")
    print(f"custom.txt — {len(lines)} lines")


def cmd_whitelist(args):
    if not os.path.exists(WHITELIST_FILE):
        open(WHITELIST_FILE, "w").close()

    with open(WHITELIST_FILE, "r") as f:
        domains = f.read().splitlines()

    if not args or args[0] == "list":
        if not domains:
            print("(whitelist.txt is empty)")
        for d in domains:
            print(d)
        return

    cmd = args[0]
    if cmd == "add" and len(args) >= 2:
        domain = args[1].strip()
        if domain not in domains:
            domains.append(domain)
    elif cmd == "del" and len(args) >= 2:
        domain = args[1].strip()
        domains = [d for d in domains if d != domain]
    else:
        eprint("Usage: whitelist list|add <domain>|del <domain>")
        return

    with open(WHITELIST_FILE, "w") as f:
        f.write("\n".join(domains) + "\n")
    print(f"whitelist.txt — {len(domains)} domains")


def cmd_help(_args):
    print("""Commands:
  list [category|all]        — list categories / items
  download                   — pick and download blocklists
  custom list|add|del|set    — manage custom overrides
  whitelist list|add|del     — manage whitelist
  help                       — this
  exit                       — quit""")


def main():
    if not os.path.isdir(BLOCKLIST_DIR):
        eprint(f"ERROR: {BLOCKLIST_DIR} does not exist. Run from repo root.")
        sys.exit(1)

    cmds = {
        "list": cmd_list,
        "download": cmd_download,
        "custom": cmd_custom,
        "whitelist": cmd_whitelist,
        "help": cmd_help,
    }

    import subprocess

    subprocess.run(["clear"])

    print("linum block control — type 'help' for commands")

    while True:
        try:
            raw = input("linum> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not raw:
            continue
        parts = raw.split()
        cmd = parts[0]
        if cmd == "exit":
            break
        fn = cmds.get(cmd)
        if fn:
            fn(parts[1:])
        else:
            eprint(f"unknown command '{cmd}'. try: help")


if __name__ == "__main__":
    main()