# base_station_main.py  --  Trivia buzzer JUDGE (Raspberry Pi Pico W)
#
# Modes: "quizbowl" (live on reset) or "jeopardy" (stand-by until activated).
# Control button GP15 -> GND:
#   quizbowl : press = reset / re-arm.
#   jeopardy : press on STAND BY = ACTIVATE; press when ARMED/LOCKED = reset.
#
# Color LEDs: the base lights the winning team's color (blue or green),
# matching the buzzer NAMEs. Onboard LED: solid = someone buzzed.

import network, socket, select, time
from machine import Pin

SSID     = "TriviaBuzzers"
PASSWORD = "buzzin123"
PORT     = 5000

MODE            = "jeopardy"    # "jeopardy" or "quizbowl"
LOCK_TIMEOUT_MS = 10000
BATT_LOW_V      = 3.60

BUTTON_PIN    = 15             # control button GP15 -> GND
BLUE_LED_PIN  = 16            # base's blue LED  -> resistor -> GND
GREEN_LED_PIN = 17            # base's green LED -> resistor -> GND

led    = Pin("LED", Pin.OUT)                 # onboard: solid = buzzed
button = Pin(BUTTON_PIN, Pin.IN, Pin.PULL_UP)

# winner NAME -> indicator LED on the base
TEAM_LEDS = {
    "blue":  Pin(BLUE_LED_PIN,  Pin.OUT),
    "green": Pin(GREEN_LED_PIN, Pin.OUT),
}

def set_team_led(name):
    for k, p in TEAM_LEDS.items():
        p.value(1 if k == name else 0)

def team_leds_off():
    for p in TEAM_LEDS.values():
        p.value(0)

# --- Wi-Fi access point ---
ap = network.WLAN(network.AP_IF)
ap.config(essid=SSID, password=PASSWORD)
ap.active(True)
while not ap.active():
    time.sleep(0.1)
print("AP ready:", ap.ifconfig())

# --- TCP server ---
srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
srv.bind(("0.0.0.0", PORT))
srv.listen(8)
srv.setblocking(False)

poller = select.poll()
poller.register(srv, select.POLLIN)

clients   = {}
phase     = "standby"
winner    = None
lock_time = 0

def send(sock, msg):
    try:
        sock.send(msg)
    except OSError:
        pass

def state_msg():
    if phase == "locked":
        return b"LOCKED " + winner.encode() + b"\n"
    if phase == "standby":
        return b"STANDBY\n"
    return b"ARMED\n"

def broadcast(msg, exclude=None):
    for c in list(clients.values()):
        if c["sock"] is not exclude:
            send(c["sock"], msg)

def battery_summary():
    if not clients:
        return
    parts = []
    for c in clients.values():
        v = c.get("batt")
        if v is None:
            parts.append("%s ?" % c["name"])
        else:
            parts.append("%s %.2f%s" % (c["name"], v, "!" if v < BATT_LOW_V else ""))
    print("Batteries:", "   ".join(parts))

def go_armed():
    global phase, winner
    phase, winner = "armed", None
    led.off()
    team_leds_off()
    broadcast(b"ARMED\n")
    print("--- ARMED (buzz now) ---")

def go_standby():
    global phase, winner
    phase, winner = "standby", None
    led.off()
    team_leds_off()
    broadcast(b"STANDBY\n")
    print("--- STAND BY ---")

def start_round():
    battery_summary()
    if MODE == "jeopardy":
        go_standby()
    else:
        go_armed()

def handle_line(fd, line):
    global phase, winner, lock_time
    c = clients.get(fd)
    if not c:
        return
    if line.startswith(b"HELLO"):
        parts = line.split()
        if len(parts) > 1:
            c["name"] = parts[1].decode()
        send(c["sock"], state_msg())
    elif line.startswith(b"BATT"):
        parts = line.split()
        if len(parts) > 1:
            try:
                v = float(parts[1])
            except ValueError:
                return
            was = c.get("batt")
            c["batt"] = v
            if v < BATT_LOW_V and (was is None or was >= BATT_LOW_V):
                print("[BATTERY LOW] %s = %.2f V" % (c["name"], v))
    elif line.startswith(b"PRESS"):
        if phase == "armed":
            phase, winner = "locked", c["name"]
            lock_time = time.ticks_ms()
            led.on()
            set_team_led(winner)
            print("WINNER:", winner)
            send(c["sock"], b"WIN\n")
            broadcast(b"LOCKED " + winner.encode() + b"\n", exclude=c["sock"])
        else:
            send(c["sock"], state_msg())

start_round()
last_btn = 0

while True:
    for sock, ev in poller.poll(20):
        if sock is srv:
            conn, addr = srv.accept()
            conn.setblocking(False)
            poller.register(conn, select.POLLIN)
            clients[conn.fileno()] = {"sock": conn, "name": str(addr[0]),
                                      "buf": b"", "batt": None}
            send(conn, state_msg())
        else:
            fd = sock.fileno()
            try:
                data = sock.recv(64)
            except OSError:
                continue
            if not data:
                poller.unregister(sock)
                clients.pop(fd, None)
                sock.close()
                continue
            c = clients.get(fd)
            if not c:
                continue
            c["buf"] += data
            while b"\n" in c["buf"]:
                line, c["buf"] = c["buf"].split(b"\n", 1)
                handle_line(fd, line.strip())

    if button.value() == 0 and time.ticks_diff(time.ticks_ms(), last_btn) > 300:
        last_btn = time.ticks_ms()
        if MODE == "jeopardy" and phase == "standby":
            go_armed()
        else:
            start_round()

    if phase == "locked" and LOCK_TIMEOUT_MS and time.ticks_diff(time.ticks_ms(), lock_time) > LOCK_TIMEOUT_MS:
        print("auto-reset after timeout")
        start_round()