vmess-domain-rotator/scripts/rotator/scoring.py

import datetime as dt
import functools

from .config import parse_timezone, to_float_or_none
from .filter import IPV4_RE, normalize_domain, resolve_field


def parse_created_time(value, formats, timezone):
    if value is None:
        return None

    s = str(value).strip()
    if not s:
        return None

    for fmt in formats:
        try:
            parsed = dt.datetime.strptime(s, fmt)
            if parsed.tzinfo is None:
                parsed = parsed.replace(tzinfo=timezone)
            return parsed.astimezone(dt.timezone.utc)
        except Exception:
            continue

    try:
        iso_text = s.replace("Z", "+00:00")
        parsed = dt.datetime.fromisoformat(iso_text)
        if parsed.tzinfo is None:
            parsed = parsed.replace(tzinfo=timezone)
        return parsed.astimezone(dt.timezone.utc)
    except Exception:
        return None


def build_lexicographic_descriptors(scoring_cfg, prefer_lower):
    out = []
    for item in scoring_cfg.get("lexicographic_fields", []):
        if isinstance(item, str):
            field_name = item.strip()
            order = "asc" if prefer_lower else "desc"
        else:
            field_name = str(item.get("field", "")).strip()
            order = str(item.get("order", "")).strip().lower()
            if not order:
                order = "asc" if prefer_lower else "desc"
        out.append({"field": field_name, "order": order})
    return out


def parse_scored_records(records, field_map, record_mapping_cfg, scoring_cfg):
    if not scoring_cfg.get("enabled", False):
        return []

    strategy = str(scoring_cfg.get("strategy", "weighted_average")).strip()
    prefer_lower = bool(scoring_cfg.get("prefer_lower", False))

    timezone = parse_timezone(record_mapping_cfg.get("created_time_timezone", "UTC"))
    time_formats = [str(x) for x in record_mapping_cfg.get("created_time_formats", [])]

    weighted_fields = scoring_cfg.get("weighted_fields", []) if strategy == "weighted_average" else []
    lex_descriptors = build_lexicographic_descriptors(scoring_cfg, prefer_lower) if strategy == "lexicographic" else []

    needed_fields = set()
    for item in weighted_fields:
        needed_fields.add(str(item.get("field", "")).strip())
    for item in lex_descriptors:
        needed_fields.add(str(item.get("field", "")).strip())
    for item in scoring_cfg.get("tie_breakers", []):
        needed_fields.add(str(item.get("field", "")).strip())
    needed_fields.discard("domain")
    needed_fields.discard("created_at")

    out = []
    for record in records:
        domain = normalize_domain(resolve_field(record, "domain", field_map))
        if not domain:
            continue

        created_raw = resolve_field(record, "created_at", field_map)
        created_at = parse_created_time(created_raw, time_formats, timezone)

        field_values = {}
        for field_name in needed_fields:
            field_values[field_name] = resolve_field(record, field_name, field_map)

        score_value = None
        scores = []
        lex_values = []

        if strategy == "weighted_average":
            total = 0.0
            total_weight = 0.0
            missing = False
            for item in weighted_fields:
                field_name = str(item.get("field", "")).strip()
                weight = float(item.get("weight"))
                raw_v = resolve_field(record, field_name, field_map)
                val = to_float_or_none(raw_v)
                scores.append(val)
                if val is None:
                    missing = True
                    continue
                field_lower = item.get("prefer_lower", prefer_lower)
                if str(item.get("order", "")).lower() == "desc":
                    field_lower = False
                elif str(item.get("order", "")).lower() == "asc":
                    field_lower = True
                effective_val = val if (bool(field_lower) == prefer_lower) else -val
                total += effective_val * weight
                total_weight += weight

            if not missing and total_weight > 0:
                score_value = total / total_weight

        if strategy == "lexicographic":
            for item in lex_descriptors:
                field_name = item["field"]
                order = item["order"]
                raw_v = resolve_field(record, field_name, field_map)
                num_v = to_float_or_none(raw_v)
                v = num_v if num_v is not None else raw_v
                lex_values.append({"field": field_name, "value": v, "order": order})
                scores.append(v)

        out.append(
            {
                "domain": domain,
                "created_at": created_at,
                "created_raw": created_raw,
                "scores": scores,
                "score_value": score_value,
                "lex_values": lex_values,
                "field_values": field_values,
                "raw_record": record,
            }
        )
    return out


def cmp_scalar(a, b, order):
    a_none = a is None
    b_none = b is None
    if a_none and b_none:
        return 0
    if a_none:
        return 1
    if b_none:
        return -1

    if isinstance(a, dt.datetime):
        a = a.timestamp()
    if isinstance(b, dt.datetime):
        b = b.timestamp()

    a_num = to_float_or_none(a)
    b_num = to_float_or_none(b)
    if a_num is not None and b_num is not None:
        if a_num < b_num:
            base = -1
        elif a_num > b_num:
            base = 1
        else:
            base = 0
    else:
        a_s = str(a).lower()
        b_s = str(b).lower()
        if a_s < b_s:
            base = -1
        elif a_s > b_s:
            base = 1
        else:
            base = 0

    return base if order == "asc" else -base


def get_sort_field_value(record, field_name):
    if field_name == "domain":
        return record.get("domain")
    if field_name == "created_at":
        return record.get("created_at")
    return record.get("field_values", {}).get(field_name)


def rank_scored_records(records, scoring_cfg):
    if not records:
        return []

    within_hours = float(scoring_cfg.get("within_hours", 0))
    strategy = str(scoring_cfg.get("strategy", "weighted_average")).strip()
    prefer_lower = bool(scoring_cfg.get("prefer_lower", False))
    tie_breakers = scoring_cfg.get("tie_breakers", [])

    if within_hours > 0:
        now = dt.datetime.now(dt.timezone.utc)
        cutoff = now - dt.timedelta(hours=within_hours)
        recent = [r for r in records if r.get("created_at") is not None and r["created_at"] >= cutoff]
        candidates = recent if recent else records
    else:
        candidates = records

    default_lex_order = "asc" if prefer_lower else "desc"

    def compare(a, b):
        if strategy == "weighted_average":
            order = "asc" if prefer_lower else "desc"
            c = cmp_scalar(a.get("score_value"), b.get("score_value"), order)
            if c != 0:
                return c
        elif strategy == "lexicographic":
            a_lex = a.get("lex_values", [])
            b_lex = b.get("lex_values", [])
            n = max(len(a_lex), len(b_lex))
            for i in range(n):
                av = a_lex[i]["value"] if i < len(a_lex) else None
                bv = b_lex[i]["value"] if i < len(b_lex) else None
                order = default_lex_order
                if i < len(a_lex) and a_lex[i].get("order"):
                    order = a_lex[i]["order"]
                c = cmp_scalar(av, bv, order)
                if c != 0:
                    return c

        for item in tie_breakers:
            field_name = str(item.get("field", "")).strip()
            order = str(item.get("order", "asc")).strip().lower()
            av = get_sort_field_value(a, field_name)
            bv = get_sort_field_value(b, field_name)
            c = cmp_scalar(av, bv, order)
            if c != 0:
                return c

        return cmp_scalar(a.get("domain"), b.get("domain"), "asc")

    return sorted(candidates, key=functools.cmp_to_key(compare))


def choose_top_candidate_domains(filtered_domains, top_n, ranked_scored):
    if ranked_scored:
        domains_by_score = [x["domain"] for x in ranked_scored]
        return domains_by_score[:top_n]
    return filtered_domains[:top_n]


def choose_domain(filtered_domains, top_n, ranked_scored, last_good="", switch_margin_pct=0.0):
    if ranked_scored:
        top_domains = choose_top_candidate_domains(filtered_domains, top_n, ranked_scored)
        if top_domains:
            best_domain = top_domains[0]
            if switch_margin_pct > 0 and last_good and best_domain != last_good:
                last_cand = next((r for r in ranked_scored if r.get("domain") == last_good), None)
                best_cand = next((r for r in ranked_scored if r.get("domain") == best_domain), None)
                if last_cand and best_cand:
                    best_score = to_float_or_none(best_cand.get("score_value"))
                    last_score = to_float_or_none(last_cand.get("score_value"))
                    if best_score is not None and last_score is not None and last_score != 0:
                        improvement_pct = abs((last_score - best_score) / last_score) * 100.0
                        if improvement_pct < switch_margin_pct:
                            top_reordered = [last_good] + [d for d in top_domains if d != last_good]
                            return last_good, top_reordered[:top_n]
            return best_domain, top_domains

    if filtered_domains:
        return filtered_domains[0], filtered_domains[:top_n]
    return None, []


def text_or_blank(value):
    if value is None:
        return ""
    return str(value).strip()


def set_if_nonempty_text(obj, key, value):
    text = text_or_blank(value)
    if text:
        obj[key] = text


def base_top_candidate(domain, source_type):
    candidate = {"domain": domain, "source_type": source_type}
    if IPV4_RE.match(domain):
        candidate["ip"] = domain
    return candidate


def maybe_resolve_field(record, field_name, field_map):
    if not isinstance(record, dict):
        return None
    if field_name not in field_map:
        return None
    return resolve_field(record, field_name, field_map)


def build_cfst_candidate(row):
    domain = row.get("domain", "")
    candidate = base_top_candidate(domain=domain, source_type="cfst_local")
    set_if_nonempty_text(candidate, "ip", row.get("ip") or domain)
    set_if_nonempty_text(candidate, "loss_rate", row.get("loss_rate"))
    set_if_nonempty_text(candidate, "avg_latency", row.get("avg_latency"))
    set_if_nonempty_text(candidate, "download_speed", row.get("download_speed"))
    set_if_nonempty_text(candidate, "region", row.get("region"))
    return candidate


def build_api_candidate(domain, record, field_map, scored_record=None):
    candidate = base_top_candidate(domain=domain, source_type="api")

    if record:
        set_if_nonempty_text(candidate, "created_raw", maybe_resolve_field(record, "created_at", field_map))
        set_if_nonempty_text(candidate, "avg_latency", maybe_resolve_field(record, "avg_latency", field_map))
        set_if_nonempty_text(candidate, "loss_rate", maybe_resolve_field(record, "avg_pkg_lost_rate", field_map))
        set_if_nonempty_text(candidate, "download_speed", maybe_resolve_field(record, "download_speed", field_map))
        location_country = text_or_blank(maybe_resolve_field(record, "location_country", field_map))
        location_city = text_or_blank(maybe_resolve_field(record, "location_city", field_map))
        region = "/".join([x for x in [location_country, location_city] if x])
        set_if_nonempty_text(candidate, "region", region)

    if scored_record:
        if scored_record.get("score_value") is not None:
            candidate["score_value"] = scored_record.get("score_value")
        scores = list(scored_record.get("scores", []))
        if scores:
            candidate["scores"] = scores
        if "created_raw" not in candidate:
            set_if_nonempty_text(candidate, "created_raw", scored_record.get("created_raw"))

    return candidate


def build_top_candidates(source_type, candidate_domains, cfst_rows=None, records=None, field_map=None, ranked_scored=None):
    if source_type == "cfst_local":
        row_map = {}
        for row in cfst_rows or []:
            domain = row.get("domain", "")
            if domain and domain not in row_map:
                row_map[domain] = row
        return [build_cfst_candidate(row_map.get(d, {"domain": d, "ip": d})) for d in candidate_domains]

    ranked_map = {}
    for item in ranked_scored or []:
        domain = item.get("domain", "")
        if domain and domain not in ranked_map:
            ranked_map[domain] = item

    record_map = {}
    for record in records or []:
        domain = normalize_domain(maybe_resolve_field(record, "domain", field_map or {}))
        if domain and domain not in record_map:
            record_map[domain] = record

    out = []
    for domain in candidate_domains:
        scored_record = ranked_map.get(domain)
        record = None
        if scored_record:
            record = scored_record.get("raw_record")
        if record is None:
            record = record_map.get(domain)
        out.append(build_api_candidate(domain, record, field_map or {}, scored_record=scored_record))
    return out