from pybricks.hubs import PrimeHub
from pybricks.tools import wait

hub = PrimeHub()


class Colors:
    RED: str = "\033[0;31m"
    GRE: str = "\033[0;32m"
    YEL: str = "\033[0;33m"
    BLU: str = "\033[0;34m"
    MAG: str = "\033[0;35m"
    CYA: str = "\033[0;36m"
    WHI: str = "\033[0;37m"
    NC: str = "\033[0m"


# ======= SINGLE SOURCE OF TRUTH =======
# 1) Curva OCV (Volt pacco 2S Li-ion -> percentuale SOC)
#    Punti ordinati per tensione, usati per un'interpolazione lineare a tratti.
OCV_POINTS: tuple[tuple[float, int], ...] = (
    (6.00, 0),  # ~3.0 V/cella: fine scarica tipica
    (6.40, 5),
    (6.60, 12),
    (6.80, 20),
    (7.00, 30),
    (7.20, 45),
    (7.40, 60),
    (7.60, 72),
    (7.80, 82),
    (8.00, 90),
    (8.19, 97),  # ~soglia "full" (LED verde) osservata su Pybricks
    (8.30, 100),  # piena carica "reale"
    (8.40, 100),  # tetto CV
)

# 2) Stati nominali in funzione della percentuale calcolata dalla curva OCV.
#    (min_percent_incluso, Nome, Colore)
BATTERY_STATES: tuple[tuple[int, str, str], ...] = (
    (90, "PIENA", Colors.BLU),
    (60, "ALTA", Colors.GRE),
    (30, "MEDIA", Colors.CYA),
    (15, "BASSA", Colors.YEL),
    (0, "CRITICA", Colors.RED),
)
# ======================================


def _clamp(x: float, lo: float, hi: float) -> float:
    return hi if x > hi else lo if x < lo else x


def _interp_soc_from_voltage(v_pack: float) -> int:
    """Interpolazione lineare a tratti sulla curva OCV."""
    v = _clamp(v_pack, OCV_POINTS[0][0], OCV_POINTS[-1][0])

    # nodo esatto?
    for vp, sp in OCV_POINTS:
        if abs(v - vp) < 1e-6:
            return int(sp)

    # trova il segmento [i, i+1] e interpola
    for i in range(len(OCV_POINTS) - 1):
        v0, s0 = OCV_POINTS[i]
        v1, s1 = OCV_POINTS[i + 1]
        if v0 <= v <= v1:
            t = (v - v0) / (v1 - v0)
            s = s0 + t * (s1 - s0)
            return int(round(s))

    return 0  # fallback (non dovrebbe accadere)


def _avg_voltage_mv(samples: int = 5, delay_ms: int = 40) -> int:
    """Media semplice per mitigare il sag sotto carico motori."""
    total = 0
    n = max(1, samples)
    for _ in range(n):
        total += hub.battery.voltage()  # mV
        wait(delay_ms)
    return total // n


def _state_from_percent(percent: int) -> tuple[str, str]:
    """Mappa la % calcolata agli stati nominali (nome+colore)."""
    p = max(0, min(100, percent))
    for p_min, name, color in BATTERY_STATES:
        if p >= p_min:
            return name, color
    # fallback al peggior stato
    return BATTERY_STATES[-1][1], BATTERY_STATES[-1][2]


def calcola_perc(v_mv: float, consider_charger: bool = True) -> int:
    """Restituisce %SOC intera da 0..100 usando la curva OCV."""
    v = v_mv / 1000.0

    if consider_charger:
        # Se il caricatore dice "full", forza 100%
        try:
            st = hub.charger.status()  # 0=off, 1=rosso, 2=verde, 3=giallo
            if st == 2:  # verde
                return 100
        except Exception:
            pass

    return _interp_soc_from_voltage(v)


def calcola_stato(percent: int) -> str:
    """Restituisce stringa colorata con il nome dello stato."""
    name, color = _state_from_percent(percent)
    return f"{color}{name}{Colors.NC}"


def print_carica(samples: int = 6) -> None:
    v_mv = _avg_voltage_mv(samples=samples)
    i_ma = hub.battery.current()
    perc = calcola_perc(v_mv, consider_charger=True)
    stato = calcola_stato(perc)

    # Info caricatore (se presente)
    charging_flag = ""
    try:
        st = hub.charger.status()
        if st == 1:
            charging_flag = "  |  In carica"
        elif st == 2:
            charging_flag = "  |  Carica completata"
    except Exception:
        pass

    print(f"Voltage: {v_mv} mV  |  Current: {i_ma} mA{charging_flag}")
    print(f"Batteria: {stato} ({perc}%)")


print_carica(samples=6)
# if __name__ == "__main__":
#     print_carica(samples=6)