#!/usr/bin/env python3
import os,time,socket,subprocess,configparser,signal
from pathlib import Path
APP_VERSION="3.1.0"
from smbus2 import SMBus
from PIL import Image
from shutil import disk_usage

OLED=0x3C; MCU=0x0D; W,H=128,32
CFG="/etc/pigway-pi-control.conf"
running=True

def stop(*_):
    global running; running=False
signal.signal(signal.SIGTERM,stop); signal.signal(signal.SIGINT,stop)

bus=SMBus(1)
FONT_PATH="/usr/local/share/pigway-pi-control/oled_font_5x7.bin"
FONT5=Path(FONT_PATH).read_bytes()
fan_level=0
prev_cpu=None
page=0; page_since=time.monotonic()
cpu_hi=mem_hi=0
net_ok=link_ok=True
ip="NO IP"
last_net=last_watch=0
watch_state=[]

def cmd(c): bus.write_byte_data(OLED,0x00,c)
def oled_init():
    # Seller-native physical 128x32 configuration.
    for c in [0xAE,0xD5,0x80,0xA8,0x1F,0xD3,0x00,0x40,0x8D,0x14,
              0x20,0x00,0xA1,0xC8,0xDA,0x02,0x81,0x7F,0xD9,0xF1,
              0xDB,0x40,0xA4,0xA6,0xAF]:
        cmd(c)

def draw_text_5x7(pix,x,y,text):
    for ch in text:
        code=ord(ch)
        if code<0 or code>255: code=ord("?")
        off=code*5
        if off+5>len(FONT5): code=ord("?"); off=code*5
        for cx in range(5):
            col=FONT5[off+cx]
            for cy in range(7):
                if col & (1<<cy):
                    xx=x+cx; yy=y+cy
                    if 0<=xx<W and 0<=yy<H: pix[xx,yy]=1
        x+=6
        if x>W-6: break

def oled_show(lines_lr, title_zh=False):
    im=Image.new("1",(W,H),0); pix=im.load()
    for row,(left,right) in enumerate(lines_lr[:4]):
        y=row*8
        draw_text_5x7(pix,0,y,left)
        if right:
            draw_text_5x7(pix,max(0,W-len(right)*6),y,right)
    data=[]
    for pg in range(4):
        for x in range(W):
            v=0
            for bit in range(8):
                if pix[x,pg*8+bit]: v|=1<<bit
            data.append(v)
    cmd(0x21);cmd(0);cmd(127);cmd(0x22);cmd(0);cmd(3)
    for i in range(0,len(data),16):
        bus.write_i2c_block_data(OLED,0x40,data[i:i+16])


def mcu(reg,val): bus.write_byte_data(MCU,reg,val)
_last_rgb=None
def rgb_off():
    global _last_rgb
    if _last_rgb != ("off",):
        mcu(0x07,0x00)  # seller closeRGB(): true OFF
        _last_rgb=("off",)

def rgb(r,g,b):
    global _last_rgb
    scale=max(0,min(100,LED_BRIGHT))/100.0
    state=(int(r*scale),int(g*scale),int(b*scale))
    if _last_rgb==state: return
    mcu(0x07,0x00)      # stop built-in effect first
    mcu(0x00,0xff)      # seller protocol: >=3 => all three LEDs
    mcu(0x01,state[0]); mcu(0x02,state[1]); mcu(0x03,state[2])
    _last_rgb=state

def set_pi_led(name,off=True):
    base=Path("/sys/class/leds")/name
    try:
        if off:
            (base/"trigger").write_text("none")
            (base/"brightness").write_text("0")
    except Exception as e:
        print(f"board LED {name}: {e}",flush=True)

def apply_dark_mode():
    if DARK_PWR: set_pi_led("PWR",True)
    if DARK_ACT: set_pi_led("ACT",True)

FAN_REG=[0x00,0x04,0x06,0x08,0x09,0x01]
FAN_NAME=["OFF","L1","L2","L3","L4","MAX"]
def fan_target(t):
    if t<FAN_T[0]: return 0
    if t<FAN_T[1]: return 1
    if t<FAN_T[2]: return 2
    if t<FAN_T[3]: return 3
    if t<FAN_T[4]: return 4
    return 5
def update_fan(t):
    global fan_level
    want=fan_target(t); old=fan_level
    if want>fan_level: fan_level=want
    elif want<fan_level:
        boundary = FAN_T[max(0,fan_level-1)] - FAN_HYST
        if t < boundary: fan_level-=1
    if old!=fan_level: mcu(0x08,FAN_REG[fan_level])

def cpu_pct():
    global prev_cpu
    p=list(map(int,Path("/proc/stat").read_text().splitlines()[0].split()[1:9]))
    idle=p[3]+p[4]; total=sum(p)
    if prev_cpu is None: prev_cpu=(idle,total); return 0.0
    di=idle-prev_cpu[0]; dt=total-prev_cpu[1]; prev_cpu=(idle,total)
    return 100.0*(dt-di)/dt if dt else 0.0
def temp_c(): return int(Path("/sys/class/thermal/thermal_zone0/temp").read_text())/1000.0
def mem_pct():
    vals={}
    for l in Path("/proc/meminfo").read_text().splitlines():
        if ":" in l:
            k,v=l.split(":",1); vals[k]=int(v.split()[0])
    return 100*(vals["MemTotal"]-vals["MemAvailable"])/vals["MemTotal"]
def disk_pct():
    d=disk_usage("/"); return 100*d.used/d.total
def get_ip():
    try:
        s=socket.socket(socket.AF_INET,socket.SOCK_DGRAM);s.connect(("1.1.1.1",53));v=s.getsockname()[0];s.close();return v
    except: return "NO IP"
def internet():
    try:
        s=socket.create_connection(("1.1.1.1",53),0.5);s.close();return True
    except: return False

def load_cfg():
    c=configparser.ConfigParser(); c.read(CFG)
    watches=[]
    if c.has_section("services"):
        for label,target in c["services"].items():
            target=target.strip()
            if target and not target.startswith("#"): watches.append(("service",label.upper(),target))
    if c.has_section("processes"):
        for label,target in c["processes"].items():
            target=target.strip()
            if target and not target.startswith("#"): watches.append(("process",label.upper(),target))
    return c,watches
cfg,watches=load_cfg()

def cf(section,key,default,cast=float):
    try: return cast(cfg.get(section,key))
    except: return default

SENSOR_INTERVAL=cf("timing","sensor_interval",1.0)
OLED_REFRESH=cf("timing","oled_refresh_interval",1.0)
PAGE_INTERVAL=cf("timing","page_interval",10.0)
NETWORK_INTERVAL=cf("timing","network_interval",5.0)
SERVICE_INTERVAL=cf("timing","service_interval",5.0)

CPU_HIGH=cf("alerts","cpu_high",90.0)
CPU_RECOVER=cf("alerts","cpu_recover",75.0)
CPU_HIGH_SEC=cf("alerts","cpu_high_seconds",30,int)
MEM_HIGH=cf("alerts","memory_high",90.0)
MEM_RECOVER=cf("alerts","memory_recover",85.0)
MEM_HIGH_SEC=cf("alerts","memory_high_seconds",30,int)
TEMP_HIGH=cf("alerts","temp_high",68.0)
TEMP_CRIT=cf("alerts","temp_critical",75.0)
DISK_HIGH=cf("alerts","disk_high",90.0)
DISK_CRIT=cf("alerts","disk_critical",95.0)

FAN_T=[cf("fan","start",50.0),cf("fan","level2",53.0),cf("fan","level3",56.0),
       cf("fan","level4",59.0),cf("fan","max",62.0)]
FAN_HYST=cf("fan","hysteresis",1.5)
ROW_Y=[cf("display","row1_y",0,int),cf("display","row2_y",8,int),cf("display","row3_y",16,int),cf("display","row4_y",24,int)]
DARK_PWR=cf("dark_mode","pi_pwr_led",True,lambda x: str(x).strip().lower() in ("off","0","false","no"))
DARK_ACT=cf("dark_mode","pi_act_led",True,lambda x: str(x).strip().lower() in ("off","0","false","no"))
LED_BRIGHT=cf("led","brightness",10,int)
LED_WARN_ON=cf("led","warning_on_ms",300,int); LED_WARN_OFF=cf("led","warning_off_ms",2700,int)
LED_CRIT_ON=cf("led","critical_on_ms",400,int); LED_CRIT_OFF=cf("led","critical_off_ms",1100,int)
LED_EMERG_ON=cf("led","emergency_on_ms",400,int); LED_EMERG_OFF=cf("led","emergency_off_ms",400,int)

def check_watch(w):
    typ,label,target=w
    if typ=="service":
        r=subprocess.run(["systemctl","is-active","--quiet",target])
    else:
        r=subprocess.run(["pgrep","-f",target],stdout=subprocess.DEVNULL)
    return r.returncode==0

def alert(title,l2="",l3="",l4="",sev=1,color=(20,16,0),blink=0):
    return {"title":title[:21],"l2":l2[:21],"l3":l3[:21],"l4":l4[:21],
            "sev":sev,"color":color,"blink":blink}

def make_alerts(cpu,temp,mem,disk):
    global cpu_hi,mem_hi
    a=[]
    cpu_hi=cpu_hi+1 if cpu>=CPU_HIGH else 0 if cpu<CPU_RECOVER else cpu_hi
    mem_hi=mem_hi+1 if mem>=MEM_HIGH else 0 if mem<MEM_RECOVER else mem_hi
    if cpu_hi>=CPU_HIGH_SEC:a.append(alert("CPU LOAD HIGH",f"LOAD {cpu:.1f}%",f"TEMP {temp:.1f}C","FOR >=30 SEC",1,(24,18,0)))
    if temp>=TEMP_CRIT:a.append(alert("TEMP CRITICAL",f"TEMP {temp:.1f}C",f"FAN {FAN_NAME[fan_level]}","CHECK COOLING",3,(32,0,0),2))
    elif temp>=TEMP_HIGH:a.append(alert("TEMP HIGH",f"TEMP {temp:.1f}C",f"FAN {FAN_NAME[fan_level]}","CHECK COOLING",1,(24,0,0)))
    if mem_hi>=MEM_HIGH_SEC:a.append(alert("MEMORY HIGH",f"USED {mem:.1f}%","CHECK PROCESS","",1,(18,0,24)))
    if disk>=DISK_CRIT:a.append(alert("DISK CRITICAL",f"USED {disk:.1f}%","FREE < 5%","CHECK STORAGE",2,(28,6,0),2))
    elif disk>=DISK_HIGH:a.append(alert("DISK SPACE LOW",f"USED {disk:.1f}%","FREE < 10%","CHECK STORAGE",1,(28,8,0)))
    if not link_ok:a.append(alert("NETWORK DOWN","NO LAN / WIFI","NO IP ADDRESS","",2,(0,8,28),1))
    elif not net_ok:a.append(alert("INTERNET DOWN","LOCAL LINK OK",ip,"",2,(0,8,28),1))
    for w,ok in zip(watches,watch_state):
        if not ok:a.append(alert("SERVICE DOWN",w[1],"CHECK / RESTART","",2,(0,24,24),1))
    a.sort(key=lambda x:x["sev"],reverse=True)
    return a

def led_for(a):
    # Color = fault object. Blink frequency = severity. All alerts blink.
    if not a:
        rgb_off(); return
    if a["sev"]>=3: on_ms,off_ms=LED_EMERG_ON,LED_EMERG_OFF
    elif a["sev"]>=2: on_ms,off_ms=LED_CRIT_ON,LED_CRIT_OFF
    else: on_ms,off_ms=LED_WARN_ON,LED_WARN_OFF
    cycle=max(1,on_ms+off_ms)
    phase=int(time.monotonic()*1000)%cycle
    rgb(*a["color"]) if phase<on_ms else rgb_off()


def render_home(cpu,temp,mem,disk):
    # 128x32 fixed columns. Both columns are LEFT-aligned.
    # Seller 5x7 font has 6px character pitch; right column starts at x=72.
    im=Image.new("1",(W,H),0); pix=im.load()
    left=[
        f"CPU {cpu:.1f}%",
        f"TMP {temp:.1f}C",
        f"DSK {disk:.1f}%",
        ("IP6 " if ":" in ip else "IP4 ")+ip,
    ]
    right=[
        f"MEM {mem:.1f}%",
        f"FAN {FAN_NAME[fan_level]}",
        f"NET {'OK' if net_ok else 'DOWN'}",
        "",
    ]
    for row in range(4):
        y=row*8
        draw_text_5x7(pix,0,y,left[row])
        if right[row]:
            draw_text_5x7(pix,72,y,right[row])
    data=[]
    for pg in range(4):
        for x in range(W):
            v=0
            for bit in range(8):
                if pix[x,pg*8+bit]: v|=1<<bit
            data.append(v)
    cmd(0x21);cmd(0);cmd(127);cmd(0x22);cmd(0);cmd(3)
    for i in range(0,len(data),16):
        bus.write_i2c_block_data(OLED,0x40,data[i:i+16])


def render_alert(a,idx,total):
    oled_show([(a["title"],f"{idx}/{total}"),
               (a["l2"],""),(a["l3"],""),(a["l4"],"")])


oled_init(); rgb_off(); apply_dark_mode()
watch_state=[check_watch(w) for w in watches]
last_sensor=last_render=0.0
last_net=last_watch=time.monotonic()
latest=(0.0,0.0,0.0,0.0)
active=[]

# Page state is a snapshot queue. Live sensor refreshes NEVER choose another page.
# Each cycle is HOME + every alert captured at the cycle boundary.
cycle_pages=["HOME"]
cycle_pos=0
page_started=time.monotonic()
current_alert_key=None

def alert_key(a):
    # Stable identity: title + monitored object/primary detail.
    return f"{a['title']}|{a['l2']}"

def find_live_alert(key, live):
    for a in live:
        if alert_key(a)==key: return a
    return None

while running:
    try:
        now=time.monotonic()


        # Independent clocks: update data/state only.
        if now-last_sensor >= SENSOR_INTERVAL:
            c=cpu_pct(); t=temp_c(); m=mem_pct(); d=disk_pct()
            latest=(c,t,m,d); update_fan(t); last_sensor=now

        if now-last_net >= NETWORK_INTERVAL:
            ip=get_ip(); link_ok=ip!="NO IP"; net_ok=internet() if link_ok else False
            last_net=now

        if now-last_watch >= SERVICE_INTERVAL:
            watch_state=[check_watch(w) for w in watches]
            last_watch=now

        c,t,m,d=latest
        active=make_alerts(c,t,m,d)

        # If the currently displayed alert truly disappeared, skip it immediately.
        if cycle_pages[cycle_pos]!="HOME":
            key=cycle_pages[cycle_pos]
            if find_live_alert(key,active) is None:
                cycle_pos += 1
                if cycle_pos>=len(cycle_pages):
                    # End of cycle: capture a fresh snapshot.
                    cycle_pages=["HOME"]+[alert_key(a) for a in active]
                    cycle_pos=0
                page_started=now

        # ONLY this timer performs normal page changes.
        if now-page_started >= PAGE_INTERVAL:
            cycle_pos += 1
            if cycle_pos>=len(cycle_pages):
                # A complete cycle ended. Snapshot all currently active alerts.
                cycle_pages=["HOME"]+[alert_key(a) for a in active]
                cycle_pos=0
            page_started=now

        # Special case: while idle on HOME with no alert pages queued, begin a new
        # alert cycle immediately when the first alert appears. HOME keeps its
        # current 10-second slot; alerts follow it and cannot be kicked out by refresh.
        if cycle_pages==["HOME"] and active:
            cycle_pages=["HOME"]+[alert_key(a) for a in active]

        # Render CURRENT page only. 1-second refresh does not affect cycle_pos.
        if now-last_render >= OLED_REFRESH:
            current=cycle_pages[cycle_pos]
            if current=="HOME":
                render_home(c,t,m,d)
            else:
                a=find_live_alert(current,active)
                if a:
                    # Number refers to this cycle's alert slot, all using same 6x8 font.
                    render_alert(a,cycle_pos,len(cycle_pages)-1)
                else:
                    render_home(c,t,m,d)
            last_render=now

        # LED follows OLED page; HOME reflects highest currently active risk.
        current=cycle_pages[cycle_pos]
        if current=="HOME":
            led_for(active[0] if active else None)
        else:
            a=find_live_alert(current,active)
            led_for(a if a else (active[0] if active else None))

        time.sleep(0.20)
    except Exception as e:
        print("loop error:",e,flush=True); time.sleep(1)

rgb_off(); mcu(0x08,0x00); bus.close()