#!/usr/bin/env python3
# SkillFishOS — enable the two disabled BC-250 CPU cores (6c/12t -> 8c/16t) at boot.
#
# Software method, no modified BIOS: writes the SMU core-presence mask (SMN 0x5A870)
# from 0x77 to 0xFF via queue-3 message 0x98 (an SMU-privileged SMN write), then warm-
# reboots once. The mask survives a warm reset, so on the second boot AGESA/PSP
# enumerate all 8 cores. A cold power-off reverts the mask, which is why this runs on
# every boot.
#
# Loop-safe & idempotent:
#   - mask already 0xFF  -> we're on the warm reboot after the unlock: do nothing.
#   - mask != 0x77       -> non-standard harvest, cores may be genuinely defective: skip.
#   - reboots ONLY after confirming the write took, so a failed write can never loop.
#
# Credit: register/primitive from github.com/rw-r-r-0644/bc250-core-unlock
import os, struct, sys, time

MASK_REG = 0x5A870
MSG_WRITE = 0x98
Q3_CMD, Q3_RSP, Q3_ARG = 0x03B10A20, 0x03B10A80, 0x03B10A88
DONE = {0x01, 0xFF, 0xFE, 0xFD, 0xFC}
CFG = "/sys/bus/pci/devices/0000:00:00.0/config"

if os.geteuid() != 0:
    sys.exit("root required")
if not os.path.exists(CFG):
    sys.exit(0)  # not a BC-250 / no PCI root complex

# --- Le due guardie decise il 16/08/2026 -------------------------------------
# Lo sblocco funziona scrivendo la maschera e poi RIAVVIANDO: e' il modo in cui
# la scheda rilegge la presenza dei core. Sono due condizioni gia' scritte
# anche nell'unita' systemd; qui si ripetono perche' prima o poi qualcuno
# lancera' questo script a mano, ed e' proprio allora che i guasti fanno male.
#
# 1. MAI da live. Nella live il riavvio arrivava pochi secondi dopo
#    l'accensione e, se capitava durante l'installazione, la faceva morire a
#    meta'. Segnalato da un utente sulla 26.06.3: "l'installer crasha e la
#    board si riavvia una volta al primo avvio" - due sintomi, un fatto solo.
try:
    with open("/proc/cmdline") as f:
        cmdline = f.read().split()
    if "boot=live" in cmdline:
        sys.stderr.write("skillfish-core-unlock: sessione live, non tocco i core "
                         "(il riavvio ucciderebbe l'installazione)\n")
        sys.exit(0)
except OSError:
    pass

# 2. Solo se l'utente l'ha chiesto. Una accensione a freddo azzera la maschera,
#    quindi la macchina si riavvia una volta a OGNI avvio da spenta: e' un
#    prezzo che ha senso pagare, ma dev'essere una scelta, non una sorpresa.
#    Il segnaposto lo crea SkillFishOS Tuner. Con --adesso si salta il
#    controllo, ed e' proprio quello che fa il Tuner quando l'utente accende
#    l'interruttore e non vuole aspettare il prossimo avvio.
SEGNAPOSTO = "/etc/skillfish/core-unlock.abilitato"
# ⚠️ --maschera deve passare di qui: e' di sola lettura e serve al postinst
#    del pacchetto proprio sulle macchine che il segnaposto NON ce l'hanno,
#    per capire se stavano gia' usando gli 8 core e non spegnerglieli sotto
#    il naso. Prima non era escluso, quindi la migrazione non scattava mai.
if ("--adesso" not in sys.argv and "--maschera" not in sys.argv
        and not os.path.exists(SEGNAPOSTO)):
    sys.stderr.write("skillfish-core-unlock: sblocco non attivato "
                     "(manca %s) - si accende da SkillFishOS Tuner\n" % SEGNAPOSTO)
    sys.exit(0)

fd = os.open(CFG, os.O_RDWR)


def rd(reg):
    os.pwrite(fd, struct.pack("<I", reg), 0xB8)
    return struct.unpack("<I", os.pread(fd, 4, 0xBC))[0]


def wr(reg, val):
    os.pwrite(fd, struct.pack("<I", reg), 0xB8)
    os.pwrite(fd, struct.pack("<I", val), 0xBC)


def send(msg, arg, budget=5.0):
    end = time.monotonic() + budget
    while rd(Q3_RSP) not in DONE and time.monotonic() < end:
        time.sleep(0.002)
    wr(Q3_RSP, 0)
    wr(Q3_ARG, arg)
    wr(Q3_ARG + 4, 0)
    wr(Q3_CMD, msg)
    end = time.monotonic() + budget
    while time.monotonic() < end:
        st = rd(Q3_RSP)
        if st in DONE:
            return st
        time.sleep(0.002)
    return None


mask = rd(MASK_REG) & 0xFF

# Sola lettura, nessun effetto: serve al postinst del pacchetto per capire se
# lo sblocco era gia' in uso su questa macchina. Sta qui e non piu' in alto
# perche' rd() ha bisogno del descrittore aperto poco sopra.
if "--maschera" in sys.argv:
    print("0x%02x" % mask)
    sys.exit(0)
if mask == 0xFF:
    sys.exit(0)  # already unlocked (warm reboot) -> keep booting with 8 cores
if mask != 0x77:
    sys.stderr.write(
        "skillfish-core-unlock: presence mask 0x%02X != 0x77 — skipping "
        "(non-standard harvest, cores may be defective)\n" % mask)
    sys.exit(0)

# Early boot: the GPU governor normally hasn't grabbed the SMU yet, but stop it just
# in case so the mailbox is free.
os.system("systemctl stop cyan-skillfish-governor.service >/dev/null 2>&1")

st = send(MSG_WRITE, MASK_REG)
if st != 0x01 or (rd(MASK_REG) & 0xFF) != 0xFF:
    sys.stderr.write("skillfish-core-unlock: SMU write failed (st=%s) — NOT rebooting\n" % st)
    os.system("systemctl start cyan-skillfish-governor.service >/dev/null 2>&1")
    sys.exit(1)  # write did not take -> no reboot, no loop

os.close(fd)
sys.stderr.write("skillfish-core-unlock: mask set to 0xFF, warm-rebooting for 8c/16t\n")
# --no-block: senza, systemctl resta in attesa che systemd esegua il
# riavvio, ma il riavvio non puo' partire finche' questo servizio non
# finisce — e questo servizio sta aspettando systemctl. Il risultato
# misurato erano 85-90 secondi di stallo a ogni avvio, con la riga
# "Call to Reboot failed: Transport endpoint is not connected" nel log,
# e un riavvio sporco che non conservava la maschera: due giri invece
# di uno. Cosi' il comando torna subito, il servizio chiude, l'avvio
# prosegue e systemd riavvia appena puo'.
os.system("systemctl --no-block reboot")
