"""根据心跳上报的显卡型号和显存，给出节点算力权重。"""

from __future__ import annotations

import re
from dataclasses import dataclass
from decimal import Decimal

HIGH = Decimal("1")
MID = Decimal("0.5")
LOW = Decimal("0.05")
NONE = Decimal("0.02")
DEFAULT_WEIGHT = LOW
GENERATION_40 = "40"
GENERATION_30 = "30"
GENERATION_10 = "10"
RECOVERY_40_DAYS = 210
RECOVERY_30_DAYS = 365
RECOVERY_10_DAYS = 540
NODE_ID_BASE = 1000
GENESIS_COUNT = 20
GENESIS_RECOVERY_DAYS = 30

_MODEL_NUMBER = re.compile(r"(RTX|GTX|GT)\s*(\d{3,4})")
_DATACENTER = ("H100", "H200", "A100", "A800", "A6000", "A40", "A30", "L40")
_INTEGRATED = ("INTEL", "UHD", "IRIS", "BASIC DISPLAY", "MICROSOFT")


@dataclass
class HardwareReport:
    gpu_model: str | None
    vram_gb: float | None
    cpu_count: int | None
    memory_gb: float | None


def display_node_id(node_index: int) -> int:
    """第 1 台显示 1001，第 2 台显示 1002。"""
    return NODE_ID_BASE + node_index


def is_genesis(node_index: int) -> bool:
    return 1 <= node_index <= GENESIS_COUNT


def displayed_population(actual: int) -> int:
    """全网台数展示时加上 1000 的基数。3 台在线显示为 1003。"""
    return NODE_ID_BASE + actual


def score_hardware(gpu_model: str | None, vram_gb: float | None) -> Decimal:
    """高端主流卡 1.0，中端 0.5，旧卡或低端独显 0.05，没有独显 0.02。"""
    name = (gpu_model or "").strip()
    upper = name.upper()
    if not name or _is_integrated(upper):
        return NONE
    if any(token in upper for token in _DATACENTER):
        return HIGH
    match = _MODEL_NUMBER.search(upper)
    if match is not None:
        return _from_model(match.group(1), int(match.group(2)))
    if vram_gb is not None and vram_gb >= 8:
        return HIGH
    if vram_gb is not None and vram_gb >= 4:
        return MID
    return LOW


def classify_generation(gpu_model: str | None) -> tuple[str, int]:
    """40 系及以上 210 天，30 系 365 天，其余包括 10 系和核显 540 天。"""
    upper = (gpu_model or "").strip().upper()
    if not upper or _is_integrated(upper):
        return GENERATION_10, RECOVERY_10_DAYS
    if any(token in upper for token in _DATACENTER):
        return GENERATION_40, RECOVERY_40_DAYS
    match = _MODEL_NUMBER.search(upper)
    if match is not None and match.group(1) == "RTX":
        number = int(match.group(2))
        if number >= 4000:
            return GENERATION_40, RECOVERY_40_DAYS
        if number >= 3000:
            return GENERATION_30, RECOVERY_30_DAYS
    return GENERATION_10, RECOVERY_10_DAYS


def _is_integrated(upper: str) -> bool:
    if any(token in upper for token in ("NVIDIA", "GEFORCE", "QUADRO")):
        return False
    if any(token in upper for token in _INTEGRATED):
        return True
    return "RADEON" in upper and "RX" not in upper


def _from_model(family: str, number: int) -> Decimal:
    if family == "RTX":
        if number >= 3060:
            return HIGH
        return MID
    if family == "GTX":
        if number >= 1060:
            return MID
        return LOW
    return LOW