#!/usr/bin/env python3
"""Собирает geosite-артефакты из data/ в несколько форматов.

Основной формат — geosite.dat (protobuf v2ray/Xray, подходит для 3x-ui).
Дополнительно: plain text, JSON, sing-box rule-set, Clash/mihomo, AdGuard.

Зависимостей нет — только стандартная библиотека. Внешние бинарники
(sing-box, mihomo) нужны лишь для компиляции .srs/.mrs и опциональны.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
DATA_DIR = ROOT / "data"

TYPE_DOMAIN = "domain"
TYPE_FULL = "full"
TYPE_KEYWORD = "keyword"
TYPE_REGEXP = "regexp"

# Значения enum Domain.Type из v2ray router protobuf.
PB_TYPE = {TYPE_KEYWORD: 0, TYPE_REGEXP: 1, TYPE_DOMAIN: 2, TYPE_FULL: 3}
# Порядок вывода: сначала суффиксы, потом точные, потом keyword/regexp.
TYPE_ORDER = {TYPE_DOMAIN: 0, TYPE_FULL: 1, TYPE_KEYWORD: 2, TYPE_REGEXP: 3}

LABEL_RE = r"(?!-)[a-z0-9_-]{1,63}(?<!-)"
DOMAIN_RE = re.compile(rf"^{LABEL_RE}(\.{LABEL_RE})+$")
ATTR_RE = re.compile(r"^[a-z0-9_-]+$")
CATEGORY_RE = re.compile(r"^[a-z0-9][a-z0-9-]*$")


class BuildError(Exception):
    """Ошибка в исходных данных — сборка должна упасть."""


@dataclass(frozen=True)
class Rule:
    type: str
    value: str
    attrs: tuple[str, ...] = ()

    @property
    def sort_key(self) -> tuple[int, str]:
        return (TYPE_ORDER[self.type], self.value)

    def __str__(self) -> str:
        base = self.value if self.type == TYPE_DOMAIN else f"{self.type}:{self.value}"
        return " ".join([base, *(f"@{a}" for a in self.attrs)])


# --------------------------------------------------------------------------- #
# Разбор исходников
# --------------------------------------------------------------------------- #


def source_path(category: str) -> Path:
    return DATA_DIR / f"{category}.txt"


def list_categories() -> list[str]:
    return sorted(p.stem for p in DATA_DIR.glob("*.txt"))


def parse_token(token: str, where: str) -> Rule:
    prefix, sep, value = token.partition(":")
    if sep and prefix.lower() in (TYPE_DOMAIN, TYPE_FULL, TYPE_KEYWORD, TYPE_REGEXP):
        rtype = prefix.lower()
    else:
        rtype, value = TYPE_DOMAIN, token

    if not value:
        raise BuildError(f"{where}: пустое значение в '{token}'")

    if rtype in (TYPE_DOMAIN, TYPE_FULL):
        try:
            value = value.encode("idna").decode("ascii")
        except UnicodeError:
            pass  # ниже отловит DOMAIN_RE
        value = value.lower().strip(".")
        if not DOMAIN_RE.match(value):
            raise BuildError(f"{where}: невалидный домен '{value}'")
    elif rtype == TYPE_KEYWORD:
        value = value.lower()
        if any(c.isspace() for c in value):
            raise BuildError(f"{where}: keyword не может содержать пробелы")
    elif rtype == TYPE_REGEXP:
        try:
            re.compile(value)
        except re.error as exc:
            raise BuildError(f"{where}: невалидный regexp '{value}': {exc}") from exc

    return Rule(rtype, value)


def parse_category(category: str) -> list[tuple[str, Rule | str]]:
    """Возвращает список ('rule', Rule) / ('include', category) в порядке файла."""
    path = source_path(category)
    if not path.is_file():
        raise BuildError(f"категория '{category}' не найдена: {path}")

    entries: list[tuple[str, Rule | str]] = []
    for lineno, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
        line = raw.split("#", 1)[0].strip()
        if not line:
            continue
        where = f"{path.name}:{lineno}"

        token, *rest = line.split()
        attrs = []
        for part in rest:
            if not part.startswith("@"):
                raise BuildError(f"{where}: ожидался атрибут '@name', получено '{part}'")
            attr = part[1:].lower()
            if not ATTR_RE.match(attr):
                raise BuildError(f"{where}: невалидный атрибут '@{attr}'")
            attrs.append(attr)

        if token.startswith("include:"):
            if attrs:
                raise BuildError(f"{where}: include не поддерживает атрибуты")
            target = token[len("include:") :]
            if not CATEGORY_RE.match(target):
                raise BuildError(f"{where}: невалидное имя категории '{target}'")
            entries.append(("include", target))
            continue

        rule = parse_token(token, where)
        entries.append(("rule", Rule(rule.type, rule.value, tuple(sorted(set(attrs))))))

    return entries


def resolve(category: str, stack: tuple[str, ...] = ()) -> list[Rule]:
    """Разворачивает include-и, убирает дубликаты, сортирует детерминированно."""
    if category in stack:
        chain = " -> ".join([*stack, category])
        raise BuildError(f"циклический include: {chain}")

    collected: dict[Rule, None] = {}
    for kind, payload in parse_category(category):
        if kind == "include":
            for rule in resolve(str(payload), (*stack, category)):
                collected.setdefault(rule, None)
        else:
            collected.setdefault(payload, None)  # type: ignore[arg-type]

    return sorted(collected, key=lambda r: r.sort_key)


# --------------------------------------------------------------------------- #
# protobuf: geosite.dat
# --------------------------------------------------------------------------- #


def _varint(value: int) -> bytes:
    out = bytearray()
    while True:
        chunk = value & 0x7F
        value >>= 7
        if value:
            out.append(chunk | 0x80)
        else:
            out.append(chunk)
            return bytes(out)


def _len_field(field: int, payload: bytes) -> bytes:
    return _varint(field << 3 | 2) + _varint(len(payload)) + payload


def _str_field(field: int, value: str) -> bytes:
    return _len_field(field, value.encode("utf-8"))


def _varint_field(field: int, value: int) -> bytes:
    return _varint(field << 3 | 0) + _varint(value)


def encode_geosite_dat(categories: dict[str, list[Rule]]) -> bytes:
    """GeoSiteList{ repeated GeoSite entry = 1 }."""
    out = bytearray()
    for name in sorted(categories):
        entry = bytearray(_str_field(1, name.upper()))  # GeoSite.country_code
        for rule in categories[name]:
            domain = bytearray()
            pb_type = PB_TYPE[rule.type]
            if pb_type:  # proto3 опускает значение по умолчанию (0 == Plain)
                domain += _varint_field(1, pb_type)
            domain += _str_field(2, rule.value)
            for attr in rule.attrs:
                # Attribute{ key = 1; bool_value = 2 }
                domain += _len_field(3, _str_field(1, attr) + _varint_field(2, 1))
            entry += _len_field(2, bytes(domain))  # GeoSite.domain
        out += _len_field(1, bytes(entry))
    return bytes(out)


def decode_geosite_dat(blob: bytes) -> dict[str, list[Rule]]:
    """Обратный разбор — используется для самопроверки собранного файла."""
    pb_to_type = {v: k for k, v in PB_TYPE.items()}

    def fields(buf: bytes):
        pos = 0
        while pos < len(buf):
            key, pos = _read_varint(buf, pos)
            field, wire = key >> 3, key & 7
            if wire == 0:
                value, pos = _read_varint(buf, pos)
                yield field, value
            elif wire == 2:
                length, pos = _read_varint(buf, pos)
                yield field, buf[pos : pos + length]
                pos += length
            else:
                raise BuildError(f"неподдерживаемый wire type {wire}")

    result: dict[str, list[Rule]] = {}
    for field, payload in fields(blob):
        if field != 1:
            continue
        name, rules = "", []
        for f, p in fields(payload):
            if f == 1:
                name = p.decode("utf-8")
            elif f == 2:
                rtype, value, attrs = TYPE_KEYWORD, "", []
                for df, dp in fields(p):
                    if df == 1:
                        rtype = pb_to_type[dp]
                    elif df == 2:
                        value = dp.decode("utf-8")
                    elif df == 3:
                        for af, ap in fields(dp):
                            if af == 1:
                                attrs.append(ap.decode("utf-8"))
                rules.append(Rule(rtype, value, tuple(attrs)))
        result[name] = rules
    return result


def _read_varint(buf: bytes, pos: int) -> tuple[int, int]:
    value = shift = 0
    while True:
        byte = buf[pos]
        pos += 1
        value |= (byte & 0x7F) << shift
        if not byte & 0x80:
            return value, pos
        shift += 7


# --------------------------------------------------------------------------- #
# Текстовые форматы
# --------------------------------------------------------------------------- #


def write(path: Path, content: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content.replace("\r\n", "\n"), encoding="utf-8", newline="\n")


def header(category: str, count: int, comment: str = "#") -> str:
    return (
        f"{comment} ipcheck-domains — категория: {category}\n"
        f"{comment} Правил: {count}. Файл генерируется автоматически, не редактируйте вручную.\n"
        f"{comment} Источник: data/{category}.txt\n"
    )


def render_plain(rules: list[Rule]) -> str:
    """Только «голые» домены — для скриптов и ручной вставки."""
    seen = {r.value for r in rules if r.type in (TYPE_DOMAIN, TYPE_FULL)}
    return "".join(f"{value}\n" for value in sorted(seen))


def render_source(category: str, rules: list[Rule]) -> str:
    return header(category, len(rules)) + "\n" + "".join(f"{r}\n" for r in rules)


def render_json(category: str, rules: list[Rule]) -> str:
    payload = {
        "category": category,
        "count": len(rules),
        "rules": [
            {"type": r.type, "value": r.value, **({"attrs": list(r.attrs)} if r.attrs else {})}
            for r in rules
        ],
    }
    return json.dumps(payload, ensure_ascii=False, indent=2) + "\n"


def render_singbox(rules: list[Rule]) -> str:
    """sing-box rule-set (source). domain: apex + domain_suffix: .apex == v2ray domain:."""
    rule: dict[str, list[str]] = {}
    for kind, values in (
        ("domain", sorted({r.value for r in rules if r.type in (TYPE_DOMAIN, TYPE_FULL)})),
        ("domain_suffix", sorted({f".{r.value}" for r in rules if r.type == TYPE_DOMAIN})),
        ("domain_keyword", sorted({r.value for r in rules if r.type == TYPE_KEYWORD})),
        ("domain_regex", sorted({r.value for r in rules if r.type == TYPE_REGEXP})),
    ):
        if values:
            rule[kind] = values
    return json.dumps({"version": 1, "rules": [rule] if rule else []}, indent=2) + "\n"


def render_clash_yaml(category: str, rules: list[Rule]) -> str:
    """rule-provider с behavior: domain — keyword/regexp этим форматом не поддерживаются."""
    lines = [header(category, len(rules)), "payload:\n"]
    for rule in rules:
        if rule.type == TYPE_DOMAIN:
            lines.append(f"  - '+.{rule.value}'\n")
        elif rule.type == TYPE_FULL:
            lines.append(f"  - '{rule.value}'\n")
    return "".join(lines)


def render_clash_list(category: str, rules: list[Rule]) -> str:
    """rule-provider с behavior: classical."""
    mapping = {
        TYPE_DOMAIN: "DOMAIN-SUFFIX",
        TYPE_FULL: "DOMAIN",
        TYPE_KEYWORD: "DOMAIN-KEYWORD",
        TYPE_REGEXP: "DOMAIN-REGEX",
    }
    body = "".join(f"{mapping[r.type]},{r.value}\n" for r in rules)
    return header(category, len(rules)) + body


def render_adguard(category: str, rules: list[Rule]) -> str:
    lines = [header(category, len(rules), comment="!")]
    for rule in rules:
        if rule.type == TYPE_DOMAIN:
            lines.append(f"||{rule.value}^\n")
        elif rule.type == TYPE_FULL:
            lines.append(f"||{rule.value}^$important\n")
    return "".join(lines)


# --------------------------------------------------------------------------- #
# Компиляция бинарных rule-set (опционально)
# --------------------------------------------------------------------------- #


def compile_binaries(out_dir: Path, categories: list[str], singbox: str | None, mihomo: str | None) -> list[str]:
    produced = []

    if singbox:
        for category in categories:
            src = out_dir / "sing-box" / f"{category}.json"
            dst = out_dir / "sing-box" / f"{category}.srs"
            subprocess.run([singbox, "rule-set", "compile", "--output", str(dst), str(src)], check=True)
            produced.append(str(dst.relative_to(out_dir)))
    else:
        print("  sing-box не найден — .srs пропущены", file=sys.stderr)

    if mihomo:
        for category in categories:
            src = out_dir / "clash" / f"{category}.yaml"
            dst = out_dir / "mihomo" / f"{category}.mrs"
            dst.parent.mkdir(parents=True, exist_ok=True)
            subprocess.run([mihomo, "convert-ruleset", "domain", "yaml", str(src), str(dst)], check=True)
            produced.append(str(dst.relative_to(out_dir)))
    else:
        print("  mihomo не найден — .mrs пропущены", file=sys.stderr)

    return produced


# --------------------------------------------------------------------------- #
# Точка входа
# --------------------------------------------------------------------------- #


def build(out_dir: Path, singbox: str | None, mihomo: str | None) -> int:
    categories = list_categories()
    if not categories:
        raise BuildError(f"в {DATA_DIR} нет ни одного .txt")

    resolved = {name: resolve(name) for name in categories}

    if out_dir.exists():
        shutil.rmtree(out_dir)
    out_dir.mkdir(parents=True)

    dat = encode_geosite_dat(resolved)
    (out_dir / "geosite.dat").write_bytes(dat)
    # Копия под отдельным именем: её можно положить рядом со штатным
    # geosite.dat и адресовать как ext:ipcheck.dat:ipcheck, ничего не перезаписывая.
    (out_dir / "ipcheck.dat").write_bytes(dat)

    # Самопроверка: файл должен читаться обратно ровно в то, что мы записали.
    decoded = decode_geosite_dat(dat)
    expected = {name.upper(): rules for name, rules in resolved.items()}
    if decoded != expected:
        raise BuildError("geosite.dat не проходит round-trip проверку")

    digest = hashlib.sha256(dat).hexdigest()
    write(out_dir / "geosite.dat.sha256sum", f"{digest}  geosite.dat\n")
    write(out_dir / "ipcheck.dat.sha256sum", f"{digest}  ipcheck.dat\n")

    for name, rules in resolved.items():
        write(out_dir / "txt" / f"{name}.txt", render_plain(rules))
        write(out_dir / "src" / f"{name}.txt", render_source(name, rules))
        write(out_dir / "json" / f"{name}.json", render_json(name, rules))
        write(out_dir / "sing-box" / f"{name}.json", render_singbox(rules))
        write(out_dir / "clash" / f"{name}.yaml", render_clash_yaml(name, rules))
        write(out_dir / "clash" / f"{name}.list", render_clash_list(name, rules))
        write(out_dir / "adguard" / f"{name}.txt", render_adguard(name, rules))

    compile_binaries(out_dir, categories, singbox, mihomo)

    epoch = os.environ.get("SOURCE_DATE_EPOCH")
    built_at = (
        datetime.fromtimestamp(int(epoch), tz=timezone.utc)
        if epoch
        else datetime.now(tz=timezone.utc)
    )
    write(
        out_dir / "metadata.json",
        json.dumps(
            {
                "built_at": built_at.replace(microsecond=0).isoformat().replace("+00:00", "Z"),
                "geosite_dat_sha256": digest,
                "geosite_dat_size": len(dat),
                "total_rules": sum(len(r) for r in resolved.values()),
                "categories": {name: len(rules) for name, rules in resolved.items()},
            },
            ensure_ascii=False,
            indent=2,
        )
        + "\n",
    )

    # Скрипт обновления едет вместе с артефактами: на сервере тогда нужен
    # ровно один базовый URL — и для .dat, и для самого апдейтера.
    updater = Path(__file__).with_name("update-geosite.sh")
    if updater.is_file():
        write(out_dir / updater.name, updater.read_text(encoding="utf-8"))

    checksums = "".join(
        f"{hashlib.sha256(path.read_bytes()).hexdigest()}  {path.relative_to(out_dir).as_posix()}\n"
        for path in sorted(out_dir.rglob("*"))
        if path.is_file()
    )
    write(out_dir / "SHA256SUMS", checksums)

    print(f"geosite.dat: {len(dat)} байт, sha256 {digest[:16]}…")
    for name in categories:
        print(f"  {name:<12} {len(resolved[name]):>4} правил")
    print(f"Готово: {out_dir}")
    return 0


def check() -> int:
    categories = list_categories()
    if not categories:
        raise BuildError(f"в {DATA_DIR} нет ни одного .txt")

    problems = []
    for name in categories:
        if not CATEGORY_RE.match(name):
            problems.append(f"{name}: имя категории должно быть [a-z0-9-]")
        entries = parse_category(name)
        seen: dict[Rule, int] = {}
        for index, (kind, payload) in enumerate(entries):
            if kind != "rule":
                continue
            rule: Rule = payload  # type: ignore[assignment]
            if rule in seen:
                problems.append(f"{name}.txt: дубликат '{rule}'")
            seen[rule] = index
        suffixes = {r.value for r in seen if r.type == TYPE_DOMAIN}
        for rule in seen:
            if rule.type != TYPE_DOMAIN:
                continue
            parent = rule.value.split(".", 1)[1] if "." in rule.value else ""
            while parent:
                if parent in suffixes:
                    problems.append(f"{name}.txt: '{rule.value}' уже покрыт '{parent}'")
                    break
                parent = parent.split(".", 1)[1] if "." in parent else ""
        resolve(name)  # ловит циклы и битые include

    for problem in problems:
        print(f"error: {problem}", file=sys.stderr)
    if problems:
        return 1
    print(f"ok: {len(categories)} категорий, замечаний нет")
    return 0


def main(argv: list[str] | None = None) -> int:
    for stream in (sys.stdout, sys.stderr):
        stream.reconfigure(encoding="utf-8", errors="replace")  # type: ignore[union-attr]

    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--out", type=Path, default=ROOT / "dist", help="каталог сборки (по умолчанию dist/)")
    parser.add_argument("--check", action="store_true", help="только проверить исходники, ничего не собирать")
    parser.add_argument("--sing-box", dest="singbox", help="путь к бинарю sing-box для сборки .srs")
    parser.add_argument("--mihomo", help="путь к бинарю mihomo для сборки .mrs")
    args = parser.parse_args(argv)

    try:
        if args.check:
            return check()
        return build(
            args.out.resolve(),
            args.singbox or shutil.which("sing-box"),
            args.mihomo or shutil.which("mihomo"),
        )
    except BuildError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())