src/utils/node_store.py

"""Unified proxy node store and health scoring.

This module keeps proxy-node concerns small and local:
- fetched subscription nodes are persisted into one unified candidate set
- runtime health is stored separately from user subscription config
- parallel requests choose nodes by cooldown + score + weighted randomness
"""

from __future__ import annotations

from dataclasses import dataclass
import json
import math
import os
import random
import threading
import time
from pathlib import Path
from typing import Any

from src.utils.logger import get_logger

logger = get_logger(__name__)

_ROOT_DIR = Path(__file__).parent.parent.parent
NODES_FILE = _ROOT_DIR / "config" / "nodes.json"
NODE_HEALTH_FILE = _ROOT_DIR / "config" / "node_health.json"
DIRECT_NODE_KEY = "__direct__"
_NODE_STORE_LOCK = threading.RLock()


@dataclass(frozen=True)
class ProxyRuntimePlan:
    """请求池运行计划：同一套流程下兼容直连、固定代理、动态代理。"""

    mode: str
    enabled_node_count: int
    request_pool_size: int
    candidate_queue_length: int
    candidate_queue_rounds: int
    node_retry_count: int
    deadline_seconds: float


@dataclass(frozen=True)
class NodeCandidateRef:
    """候选队列只返回节点标识和调度元信息，配置由单独接口读取。"""

    node_key: str
    index: int
    mode: str
    score: float = 0.0
    name: str = ""
    cooldown_until: float = 0.0
    selection_reason: str = ""


def _default_json_value(default: Any) -> Any:
    if isinstance(default, dict):
        return dict(default)
    if isinstance(default, list):
        return list(default)
    return default


def _read_json(path: Path, default: Any) -> Any:
    with _NODE_STORE_LOCK:
        try:
            if not path.exists():
                return _default_json_value(default)
            with open(path, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception as e:
            backup = path.with_suffix(path.suffix + ".bak")
            if backup.exists():
                try:
                    with open(backup, "r", encoding="utf-8") as f:
                        data = json.load(f)
                    logger.warning(f"读取节点存储失败，已使用备份: {path} {e}")
                    return data
                except Exception as backup_error:
                    logger.warning(f"读取节点存储备份也失败: {backup} {backup_error}")
            logger.warning(f"读取节点存储失败: {path} {e}")
            return _default_json_value(default)


def _dump_json_file(path: Path, data: Any) -> None:
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
        f.write("\n")
        f.flush()
        os.fsync(f.fileno())


def _write_json(path: Path, data: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    unique = f"{os.getpid()}.{threading.get_ident()}.{time.time_ns()}"
    tmp = path.with_name(f".{path.name}.{unique}.tmp")
    backup = path.with_suffix(path.suffix + ".bak")
    backup_tmp = path.with_name(f".{path.name}.{unique}.bak.tmp")
    with _NODE_STORE_LOCK:
        try:
            _dump_json_file(tmp, data)
            os.replace(tmp, path)
            try:
                _dump_json_file(backup_tmp, data)
                os.replace(backup_tmp, backup)
            except Exception as backup_error:
                logger.debug(f"写入节点存储备份失败: {backup} {backup_error}")
        finally:
            for leftover in (tmp, backup_tmp):
                try:
                    if leftover.exists():
                        leftover.unlink()
                except Exception:
                    pass


def node_key(raw_uri: str) -> str:
    return raw_uri.strip()


def _default_health() -> dict[str, Any]:
    return {
        "success_count": 0,
        "fail_count": 0,
        "consecutive_failures": 0,
        "first_chunk_success_count": 0,
        "stream_complete_count": 0,
        "stream_fail_count": 0,
        "stream_stall_count": 0,
        "avg_first_chunk_ms": 0.0,
        "avg_stream_stall_gap_ms": 0.0,
        "last_success_at": 0,
        "last_fail_at": 0,
        "last_stream_stall_at": 0,
        "last_stream_fail_at": 0,
        "last_stream_complete_at": 0,
        "cooldown_until": 0,
        "last_error": "",
        "last_stream_error": "",
    }


def load_nodes() -> list[dict[str, Any]]:
    data = _read_json(NODES_FILE, {"nodes": []})
    nodes = data.get("nodes", []) if isinstance(data, dict) else []
    return [n for n in nodes if isinstance(n, dict) and str(n.get("raw_uri", "")).strip()]


def _is_truthy_enabled(value: Any) -> bool:
    if isinstance(value, bool):
        return value
    if value is None:
        return False
    return str(value).strip().lower() in {"1", "true", "yes", "on", "enabled", "启用"}


def load_enabled_nodes() -> list[dict[str, Any]]:
    """返回用户明确启用的节点；未启用节点只保留在配置列表中，不参与调度。"""
    return [node for node in load_nodes() if _is_truthy_enabled(node.get("enabled"))]


def save_nodes(nodes: list[dict[str, Any]], source: str = "subscriptions") -> None:
    seen: set[str] = set()
    cleaned: list[dict[str, Any]] = []
    now = int(time.time())
    for node in nodes:
        raw_uri = str(node.get("raw_uri") or "").strip()
        if not raw_uri or raw_uri in seen:
            continue
        seen.add(raw_uri)
        item = dict(node)
        item["raw_uri"] = raw_uri
        item.setdefault("name", raw_uri[:40])
        item.setdefault("source", source)
        item.setdefault("enabled", False)
        item["updated_at"] = now
        cleaned.append(item)
    _write_json(NODES_FILE, {"updated_at": now, "count": len(cleaned), "nodes": cleaned})
    sync_health_for_nodes(cleaned)


def upsert_nodes(nodes: list[dict[str, Any]], source: str = "manual") -> tuple[list[dict[str, Any]], int, int]:
    """Merge nodes into the unified store by raw_uri.

    Returns (all_nodes, added_count, updated_count).
    """
    existing = load_nodes()
    by_uri: dict[str, dict[str, Any]] = {str(n.get("raw_uri") or "").strip(): dict(n) for n in existing}
    now = int(time.time())
    added = 0
    updated = 0

    for node in nodes:
        raw_uri = str(node.get("raw_uri") or "").strip()
        if not raw_uri:
            continue
        item = dict(node)
        item["raw_uri"] = raw_uri
        item.setdefault("name", raw_uri[:40])
        item.setdefault("source", source)
        item.setdefault("enabled", False)
        item["updated_at"] = now
        if raw_uri in by_uri:
            old = by_uri[raw_uri]
            old.update({k: v for k, v in item.items() if v not in (None, "")})
            old["updated_at"] = now
            by_uri[raw_uri] = old
            updated += 1
        else:
            by_uri[raw_uri] = item
            added += 1

    merged = list(by_uri.values())
    _write_json(NODES_FILE, {"updated_at": now, "count": len(merged), "nodes": merged})
    sync_health_for_nodes(merged)
    return merged, added, updated


def replace_node(raw_uri: str, updates: dict[str, Any]) -> dict[str, Any]:
    raw_uri = raw_uri.strip()
    nodes = load_nodes()
    now = int(time.time())
    for idx, node in enumerate(nodes):
        if str(node.get("raw_uri") or "").strip() != raw_uri:
            continue
        item = dict(node)
        for key in ("name", "server", "port", "type", "usable_as_proxy", "raw_uri", "source", "subscription_name", "enabled"):
            if key in updates:
                item[key] = updates[key]
        item["raw_uri"] = str(item.get("raw_uri") or raw_uri).strip()
        item["updated_at"] = now
        nodes[idx] = item
        _write_json(NODES_FILE, {"updated_at": now, "count": len(nodes), "nodes": nodes})
        health = load_health()
        old_key = node_key(raw_uri)
        new_key = node_key(str(item.get("raw_uri") or raw_uri))
        if old_key != new_key and old_key in health:
            health[new_key] = health.pop(old_key)
            save_health(health)
        sync_health_for_nodes(nodes)
        return item
    raise KeyError("node not found")


def set_nodes_enabled(raw_uris: list[str], enabled: bool) -> tuple[int, int, int, int]:
    """Batch-toggle node enabled state with a single read/write.

    Returns (matched_count, updated_count, enabled_count_after, total_count).
    """
    targets = {u.strip() for u in raw_uris if u.strip()}
    if not targets:
        nodes = load_nodes()
        return 0, 0, sum(1 for node in nodes if _is_truthy_enabled(node.get("enabled"))), len(nodes)

    nodes = load_nodes()
    now = int(time.time())
    matched = 0
    updated = 0
    changed = False

    for idx, node in enumerate(nodes):
        raw_uri = str(node.get("raw_uri") or "").strip()
        if raw_uri not in targets:
            continue
        matched += 1
        if _is_truthy_enabled(node.get("enabled")) == enabled:
            continue
        item = dict(node)
        item["enabled"] = enabled
        item["updated_at"] = now
        nodes[idx] = item
        updated += 1
        changed = True

    if changed:
        _write_json(NODES_FILE, {"updated_at": now, "count": len(nodes), "nodes": nodes})

    enabled_count = sum(1 for node in nodes if _is_truthy_enabled(node.get("enabled")))
    return matched, updated, enabled_count, len(nodes)


def delete_nodes(raw_uris: list[str]) -> int:
    targets = {u.strip() for u in raw_uris if u.strip()}
    if not targets:
        return 0
    data = _read_json(NODES_FILE, {"nodes": []})
    raw_nodes = data.get("nodes", []) if isinstance(data, dict) else []
    nodes = [n for n in raw_nodes if isinstance(n, dict) and str(n.get("raw_uri", "")).strip()]
    kept = [n for n in nodes if str(n.get("raw_uri") or "").strip() not in targets]
    deleted = len(nodes) - len(kept)
    if deleted:
        now = int(time.time())
        _write_json(NODES_FILE, {"updated_at": now, "count": len(kept), "nodes": kept})
        health = load_health()
        for raw_uri in targets:
            health.pop(node_key(raw_uri), None)
        save_health(health)
    return deleted


def clear_nodes() -> int:
    nodes = load_nodes()
    now = int(time.time())
    _write_json(NODES_FILE, {"updated_at": now, "count": 0, "nodes": []})
    save_health({})
    return len(nodes)


def load_health() -> dict[str, dict[str, Any]]:
    data = _read_json(NODE_HEALTH_FILE, {})
    if not isinstance(data, dict):
        return {}
    return {str(k): v for k, v in data.items() if isinstance(v, dict)}


def save_health(health: dict[str, dict[str, Any]]) -> None:
    _write_json(NODE_HEALTH_FILE, health)


def sync_health_for_nodes(nodes: list[dict[str, Any]] | None = None) -> dict[str, dict[str, Any]]:
    """让运行状态列表与节点配置列表同步：新增初始化，删除清理。"""
    nodes = load_nodes() if nodes is None else nodes
    keys = {node_key(str(n.get("raw_uri") or "")) for n in nodes if str(n.get("raw_uri") or "").strip()}
    health = load_health()
    changed = False
    for key in keys:
        if key not in health:
            health[key] = _default_health()
            changed = True
        else:
            before = dict(health[key])
            _health_for(health, key)
            changed = changed or before != health[key]
    for key in list(health.keys()):
        if key not in keys:
            health.pop(key, None)
            changed = True
    if changed:
        save_health(health)
    return health


def _health_for(health: dict[str, dict[str, Any]], raw_uri: str) -> dict[str, Any]:
    key = node_key(raw_uri)
    item = health.setdefault(key, _default_health())
    for k, v in _default_health().items():
        item.setdefault(k, v)
    return item


def node_score(node: dict[str, Any], health: dict[str, Any], now: float | None = None) -> float:
    now = now or time.time()
    cooldown_until = float(health.get("cooldown_until") or 0)
    if cooldown_until > now:
        return float("-inf")

    success_count = int(health.get("success_count") or 0)
    fail_count = int(health.get("fail_count") or 0)
    consecutive_failures = int(health.get("consecutive_failures") or 0)
    first_chunk_success_count = int(health.get("first_chunk_success_count") or 0)
    stream_fail_count = int(health.get("stream_fail_count") or 0)
    stream_stall_count = int(health.get("stream_stall_count") or 0)
    stream_complete_count = int(health.get("stream_complete_count") or 0)
    avg_first_chunk_ms = float(health.get("avg_first_chunk_ms") or 0)
    avg_stream_stall_gap_ms = float(health.get("avg_stream_stall_gap_ms") or 0)
    last_success_at = float(health.get("last_success_at") or 0)
    last_fail_at = float(health.get("last_fail_at") or 0)
    last_seen = max(last_success_at, last_fail_at)

    score = 100.0
    # 完整流式成功是最高价值信号；首包成功但后续失败次之；首包前失败进入冷却并强惩罚。
    score += min(stream_complete_count, 100) * 12.0
    score += min(success_count, 100) * 4.0
    score += min(first_chunk_success_count, 100) * 3.0
    score -= min(stream_fail_count, 100) * 1.5
    score -= min(stream_stall_count, 100) * 12.0
    score -= min(fail_count, 100) * 8.0
    score -= consecutive_failures * 35.0
    if avg_first_chunk_ms > 0:
        score -= min(avg_first_chunk_ms / 1000.0, 30.0)
    if avg_stream_stall_gap_ms > 0:
        score -= min(avg_stream_stall_gap_ms / 1000.0, 60.0)
    if not last_seen:
        score += 20.0
    elif now - last_seen > 3600:
        score += 10.0
    if str(node.get("usable_as_proxy", "")).lower() == "true" or node.get("usable_as_proxy") is True:
        score += 2.0
    return score


def _cfg_int(cfg: dict[str, Any], key: str, fallback_key: str | None, default: int) -> int:
    value = cfg.get(key)
    if value is None and fallback_key:
        value = cfg.get(fallback_key)
    try:
        return int(value if value is not None else default)
    except Exception:
        return default


def resolve_proxy_runtime_plan(cfg: dict[str, Any], enabled_node_count: int) -> ProxyRuntimePlan:
    """根据实际节点数量自适应请求池参数，但保持统一请求池流程。"""
    max_size = max(1, _cfg_int(cfg, "parallel_pool_max_size", None, 12))
    configured_pool_size = max(1, min(_cfg_int(cfg, "parallel_pool_size", None, 4), max_size))
    configured_queue_length = max(1, _cfg_int(cfg, "candidate_queue_length", None, 80))
    configured_rounds = max(0, _cfg_int(cfg, "candidate_queue_rounds", None, 0))
    node_retry_count = max(0, _cfg_int(cfg, "node_retry_count", None, 0))
    try:
        deadline_seconds = max(0.0, float(cfg.get("request_pool_deadline_seconds", 0) or 0))
    except Exception:
        deadline_seconds = 0.0
    if enabled_node_count <= 0:
        return ProxyRuntimePlan("direct", 0, 1, 1, 1, node_retry_count, deadline_seconds)
    if enabled_node_count == 1:
        rounds = configured_rounds if configured_rounds > 0 else 1
        return ProxyRuntimePlan("fixed", 1, 1, 1, rounds, node_retry_count, deadline_seconds)

    queue_length = max(1, min(configured_queue_length, enabled_node_count))
    pool_size = max(1, min(configured_pool_size, queue_length, enabled_node_count))
    return ProxyRuntimePlan("dynamic", enabled_node_count, pool_size, queue_length, configured_rounds, node_retry_count, deadline_seconds)


def get_node_config(key: str) -> dict[str, Any] | None:
    if key == DIRECT_NODE_KEY:
        return {"node_key": DIRECT_NODE_KEY, "raw_uri": "", "name": "直连", "mode": "direct"}
    for node in load_enabled_nodes():
        raw_uri = str(node.get("raw_uri") or "").strip()
        if node_key(raw_uri) == key:
            item = dict(node)
            item["node_key"] = key
            item.setdefault("mode", "proxy")
            return item
    return None


def _weighted_without_replacement(items: list[tuple[float, int, dict[str, Any]]], limit: int) -> list[tuple[float, int, dict[str, Any]]]:
    selected: list[tuple[float, int, dict[str, Any]]] = []
    pool = list(items)
    while pool and len(selected) < limit:
        weights = [max(1.0, item[0] + 120.0) if not math.isinf(item[0]) else 1.0 for item in pool]
        chosen = random.choices(pool, weights=weights, k=1)[0]
        pool.remove(chosen)
        selected.append(chosen)
    return selected


def get_candidate_queue(cfg: dict[str, Any], plan: ProxyRuntimePlan | None = None) -> list[NodeCandidateRef]:
    nodes = load_enabled_nodes()
    plan = plan or resolve_proxy_runtime_plan(cfg, len(nodes))
    if plan.mode == "direct":
        return [NodeCandidateRef(DIRECT_NODE_KEY, 0, "direct", 0.0, "直连", 0.0, "direct")]

    if not nodes:
        return []

    health = sync_health_for_nodes(nodes)
    now = time.time()
    scored: list[tuple[float, int, dict[str, Any]]] = []
    cooled: list[tuple[float, int, dict[str, Any]]] = []

    def item_key(item: tuple[float, int, dict[str, Any]]) -> str:
        return node_key(str(item[2].get("raw_uri") or ""))

    def item_health(item: tuple[float, int, dict[str, Any]]) -> dict[str, Any]:
        return _health_for(health, str(item[2].get("raw_uri") or ""))

    for idx, node in enumerate(nodes):
        raw_uri = str(node.get("raw_uri") or "")
        h = _health_for(health, raw_uri)
        score = node_score(node, h, now)
        item = (score, idx, node)
        if math.isinf(score) and score < 0:
            cooled.append(item)
        else:
            scored.append(item)

    complete_winners: list[tuple[float, int, dict[str, Any]]] = []
    interrupted_winners: list[tuple[float, int, dict[str, Any]]] = []
    regular_available: list[tuple[float, int, dict[str, Any]]] = []
    for item in scored:
        h = item_health(item)
        stream_complete_count = int(h.get("stream_complete_count") or 0)
        stream_fail_count = int(h.get("stream_fail_count") or 0)
        first_chunk_success_count = int(h.get("first_chunk_success_count") or 0)
        last_complete_at = float(h.get("last_stream_complete_at") or 0)
        last_stream_fail_at = float(h.get("last_stream_fail_at") or 0)
        if stream_complete_count > 0 and last_complete_at >= last_stream_fail_at:
            complete_winners.append(item)
        elif stream_fail_count > 0 or first_chunk_success_count > 0:
            interrupted_winners.append(item)
        else:
            regular_available.append(item)

    complete_winners.sort(
        key=lambda x: (x[0], float(item_health(x).get("last_stream_complete_at") or 0)),
        reverse=True,
    )
    interrupted_winners.sort(
        key=lambda x: (x[0], float(item_health(x).get("last_stream_fail_at") or 0)),
        reverse=True,
    )

    selected: list[tuple[float, int, dict[str, Any]]] = []
    selected_keys: set[str] = set()
    selected_reasons: dict[str, str] = {}

    def append_ordered(items: list[tuple[float, int, dict[str, Any]]], reason: str) -> None:
        for item in items:
            if len(selected) >= plan.candidate_queue_length:
                return
            key = item_key(item)
            if key in selected_keys:
                continue
            selected.append(item)
            selected_keys.add(key)
            selected_reasons[key] = reason

    append_ordered(complete_winners, "winner_complete")
    append_ordered(interrupted_winners, "winner_interrupted")

    remaining_regular = [item for item in regular_available if item_key(item) not in selected_keys]
    weighted_regular = _weighted_without_replacement(remaining_regular, plan.candidate_queue_length - len(selected))
    append_ordered(weighted_regular, "weighted_score")

    if len(selected) < plan.candidate_queue_length:
        cooled.sort(key=lambda x: float(_health_for(health, str(x[2].get("raw_uri") or "")).get("cooldown_until") or 0))
        append_ordered(cooled[:plan.candidate_queue_length - len(selected)], "cooldown_near_end")

    refs: list[NodeCandidateRef] = []
    for score, idx, node in selected:
        raw_uri = str(node.get("raw_uri") or "").strip()
        h = _health_for(health, raw_uri)
        refs.append(NodeCandidateRef(
            node_key(raw_uri),
            idx,
            "proxy",
            score,
            str(node.get("name") or raw_uri[:40] or f"node-{idx+1}"),
            float(h.get("cooldown_until") or 0),
            selected_reasons.get(node_key(raw_uri), "weighted_score"),
        ))
    return refs


def select_nodes_for_parallel(cfg: dict[str, Any], requested_count: int) -> list[tuple[int, dict[str, Any]]]:
    cfg = dict(cfg)
    cfg["candidate_queue_length"] = requested_count
    plan = resolve_proxy_runtime_plan(cfg, len(load_enabled_nodes()))
    selected: list[tuple[int, dict[str, Any]]] = []
    for ref in get_candidate_queue(cfg, plan):
        node = get_node_config(ref.node_key)
        if node is None:
            continue
        selected.append((ref.index, node))
    return selected


def record_node_success(node: dict[str, Any], first_chunk_ms: float) -> None:
    raw_uri = str(node.get("raw_uri") or "").strip()
    if not raw_uri:
        return
    health = load_health()
    item = _health_for(health, raw_uri)
    item["success_count"] = int(item.get("success_count") or 0) + 1
    item["first_chunk_success_count"] = int(item.get("first_chunk_success_count") or 0) + 1
    item["consecutive_failures"] = 0
    item["last_success_at"] = int(time.time())
    item["cooldown_until"] = 0
    old_avg = float(item.get("avg_first_chunk_ms") or 0)
    item["avg_first_chunk_ms"] = first_chunk_ms if old_avg <= 0 else old_avg * 0.7 + first_chunk_ms * 0.3
    save_health(health)


def record_node_failure(node: dict[str, Any], error: Exception | str) -> None:
    raw_uri = str(node.get("raw_uri") or "").strip()
    if not raw_uri:
        return
    health = load_health()
    item = _health_for(health, raw_uri)
    item["fail_count"] = int(item.get("fail_count") or 0) + 1
    item["consecutive_failures"] = int(item.get("consecutive_failures") or 0) + 1
    now = int(time.time())
    item["last_fail_at"] = now
    failures = max(1, int(item["consecutive_failures"]))
    cooldown = min(1800, 30 * (2 ** min(failures - 1, 6)))
    item["cooldown_until"] = now + cooldown
    item["last_error"] = str(error)[:500]
    save_health(health)


def record_node_stream_complete(node: dict[str, Any]) -> None:
    raw_uri = str(node.get("raw_uri") or "").strip()
    if not raw_uri:
        return
    health = load_health()
    item = _health_for(health, raw_uri)
    item["stream_complete_count"] = int(item.get("stream_complete_count") or 0) + 1
    item["success_count"] = int(item.get("success_count") or 0) + 1
    item["consecutive_failures"] = 0
    now = int(time.time())
    item["last_stream_complete_at"] = now
    item["last_success_at"] = now
    item["cooldown_until"] = 0
    save_health(health)


def record_node_stream_failure(node: dict[str, Any], error: Exception | str) -> None:
    """首包后流中断：保留首包成功价值，只轻微降低稳定性评分，不进入冷却通道。"""
    raw_uri = str(node.get("raw_uri") or "").strip()
    if not raw_uri:
        return
    health = load_health()
    item = _health_for(health, raw_uri)
    item["stream_fail_count"] = int(item.get("stream_fail_count") or 0) + 1
    item["consecutive_failures"] = 0
    now = int(time.time())
    item["last_stream_fail_at"] = now
    item["last_success_at"] = max(int(item.get("last_success_at") or 0), now)
    item["cooldown_until"] = 0
    item["last_stream_error"] = str(error)[:500]
    save_health(health)


def record_node_stream_stall(node: dict[str, Any], gap_ms: float, error: Exception | str) -> None:
    """winner 首包后 raw chunk 长时间停顿：明显降低流式稳定性评分，但不禁用节点。"""
    raw_uri = str(node.get("raw_uri") or "").strip()
    if not raw_uri:
        return
    health = load_health()
    item = _health_for(health, raw_uri)
    item["stream_stall_count"] = int(item.get("stream_stall_count") or 0) + 1
    item["stream_fail_count"] = int(item.get("stream_fail_count") or 0) + 1
    item["consecutive_failures"] = 0
    now = int(time.time())
    item["last_stream_stall_at"] = now
    item["last_stream_fail_at"] = now
    item["last_success_at"] = max(int(item.get("last_success_at") or 0), now)
    item["cooldown_until"] = 0
    old_avg = float(item.get("avg_stream_stall_gap_ms") or 0)
    item["avg_stream_stall_gap_ms"] = gap_ms if old_avg <= 0 else old_avg * 0.7 + gap_ms * 0.3
    item["last_stream_error"] = str(error)[:500]
    save_health(health)