#!/usr/bin/env python3
import os
import sys
import re
import time
import psutil
import signal
import pyudev
import threading
import glob
import subprocess
import numpy as np
from PIL import Image, ImageDraw, ImageFont, ImageOps

# --- Настройки ---
# ASM1166 port map: 3, 4, 5, 6, 2, 1
PORT_MAP = {3: 0, 4: 1, 5: 2, 6: 3, 2: 4, 1: 5}
LCD_DIR = "/etc/lcd"
FB_PATH = "/dev/fb_lcd"
PWM_CHIP = "/sys/class/pwm/pwmchip0"
PWM_PATH = PWM_CHIP + "/pwm0"
W, H = 128, 64

shared_slots = [None] * 6 
shared_health = {}
shared_temps = {'cpu': 0, 'disk': 0}
lock = threading.Lock()
active_devs = {}

try:
	font = ImageFont.truetype("/usr/share/fonts/truetype/liberation/LiberationMono-Regular.ttf", 9)
except:
	font = ImageFont.load_default()

# --- Вспомогательные функции ---
def load_res(name, size=None):
	p = os.path.join(LCD_DIR, name)
	if os.path.exists(p):
		img = Image.open(p).convert('L')
		if size: img = img.resize(size)
		return ImageOps.invert(img).convert('1')
	return Image.new('1', size if size else (8,8), 0)

def send_fb(img):
	try:
		raw = np.array(img.convert('L'))
		raw = (raw > 0).astype(np.uint8) * 255
		rgb = np.dstack([raw, raw, raw, np.zeros_like(raw)])
		with open(FB_PATH, "wb") as f:
			f.write(rgb.tobytes())
	except:
		pass

def setup_pwm():
	try:
		if not os.path.exists(PWM_PATH):
			with open(f"{PWM_CHIP}/export", "w") as f:
				f.write("0")
		time.sleep(0.2)
		with open(f"{PWM_PATH}/period", "w") as f:
			f.write("40000") # 25kHz
		with open(f"{PWM_PATH}/duty_cycle", "w") as f:
			f.write("40000")
		with open(f"{PWM_PATH}/enable", "w") as f:
			f.write("1")
	except:
		pass

def get_cpu_t():
	try:
		for p in glob.glob("/sys/class/thermal/thermal_zone*/"):
			with open(p + "type") as f:
				if "cpu" in f.read():
					with open(p + "temp") as f: return int(f.read())/1000
	except: pass
	return 0

# --- Поток 1: Сканер дисков и Активности ---
def disk_scanner_worker():
	global shared_slots
	last_io = {}

	while True:
		new_slots = [None] * 6
		# 1. Быстро находим диски только на контроллере ASM1166
		for dev_path in glob.glob("/sys/block/sd*"):
			dev = os.path.basename(dev_path)
			
			# ПРОВЕРКА 1: Если диска нет в /dev, даже не пытаемся лезть в /sys (защита от фриза)
			if not os.path.exists(f"/dev/{dev}"):
				continue
				
			try:
				# ПРОВЕРКА 2: Читаем путь порта (ataX)
				# os.readlink обычно не виснет, если /dev/sdX на месте
				real_p = os.readlink(dev_path)
				if "6000000.pcie" not in real_p: continue # Игнорируем USB
				
				# Ищем номер порта
				import re
				m = re.search(r'ata(\d+)', real_p)
				if m:
					port_num = int(m.group(1))
					idx = PORT_MAP.get(port_num)
					if idx is not None:
						# Читаем активность
						with open(f"{dev_path}/stat", 'r') as f:
							cur = int(f.read().split()[9])
						
						is_act = cur > last_io.get(dev, 0)
						last_io[dev] = cur
						
						new_slots[idx] = {
							'dev': dev, 'act': is_act, 
							'ok': shared_health.get(dev, True)
						}
			except: continue

		# 2. Температуры дисков (hwmon)
		d_t = 0
		for n in glob.glob("/sys/class/hwmon/hwmon*/temp1_input"):
			try:
				with open(n) as f: d_t = max(d_t, int(f.read())/1000)
			except: pass

		with lock:
			shared_slots = new_slots
			shared_temps['disk'] = d_t
		time.sleep(0.5)

# --- Поток: SMART ---
def smart_worker():
	while True:
		# Опрашиваем только то, что сейчас есть в системе
		devs = [d for d in os.listdir('/dev') if d.startswith('sd') and len(d) == 3]
		for d in devs:
			try:
				# Жесткий таймаут 2 сек, чтобы не вешать контроллер
				res = subprocess.run(['smartctl', '-H', f'/dev/{d}'], 
									 capture_output=True, text=True, timeout=2)
				shared_health[d] = "PASSED" in res.stdout
			except: pass
		time.sleep(60)

# --- Поток 3: Анимация загрузки ---
def boot_anim():
	i = 0
	bg = load_res("boot.png", (W, H))

	while True:
		tmp = bg.copy()
		draw = ImageDraw.Draw(tmp)
		bx, by = 34, 52
		draw.rectangle((bx, by, bx+60, by+6), outline=1)
		shift = (i % 8) * 7
		draw.rectangle((bx+2+shift, by+2, bx+5+shift, by+4), fill=1)
		send_fb(tmp)
		i += 1
		time.sleep(0.2)
		try:
			if "running" in subprocess.check_output(["systemctl", "is-system-running"], text=True): break
		except: pass

# --- Основной цикл (Main) ---
setup_pwm()

img_shut = load_res("shutdown.png")
icons = {
	"ok": load_res("disk_ok.png"),
	"act": load_res("disk_active.png"),
	"err": load_res("disk_error.png")
}

# Запуск потоков
t_boot = threading.Thread(target=boot_anim)
t_boot.start()
t_boot.join()

threading.Thread(target=disk_scanner_worker, daemon=True).start()
threading.Thread(target=smart_worker, daemon=True).start()
signal.signal(signal.SIGTERM, lambda s, f: (send_fb(img_shut), sys.exit(0)))

# --- Основной цикл ---
print("Main thread started.")
blink = True
fan_on = False
while True:
	try:
		img = Image.new('1', (W, H), 0)
		draw = ImageDraw.Draw(img)
		blink = not blink
		
		# 1. Диски (Берем из памяти)
		with lock:
			current_view = list(shared_slots)
			t_d = shared_temps['disk']

		for i in range(6):
			slot = current_view[i]
			x = i * 12 + 2
			if slot:
				if not slot['ok']:
					ic = icons['err'] if blink else icons['ok']
				elif slot['act']:
					ic = icons['act']
				else:
					ic = icons['ok']
				img.paste(ic, (x, 2))
			else:
				draw.line((x + 2, 6, x + 8, 6), fill=1)

		draw.text((82, 1), time.strftime("%H:%M:%S"), font=font, fill=1)
		draw.line((0, 12, 128, 12), fill=1)

		# 2. Метрики (CPU/RAM/SYS)
		m = [("CPU", psutil.cpu_percent()), ("RAM", psutil.virtual_memory().percent),
			 ("SWP", psutil.swap_memory().percent), ("SYS", psutil.disk_usage('/').percent)]
		for i, (l, v) in enumerate(m):
			y = 15 + (i * 9)
			draw.text((2, y-1), l, font=font, fill=1)
			draw.rectangle((28, y+1, 98, y+6), outline=1)
			draw.rectangle((30, y+3, 30 + int(66 * v/100), y+4), fill=1)
			draw.text((102, y-1), f"{int(v)}%", font=font, fill=1)

		# 3. Температура и Вентилятор
		# Читаем CPU из sysfs
		t_c = get_cpu_t()
		target = max(t_c, t_d)
		if target > 40: fan_on = True
		elif target < 37: fan_on = False

		p_pct, duty = 0, 40000
		if fan_on:
			ratio = min(1.0, max(0, (target - 40) / 15))
			duty = int(24000 * (1.0 - ratio))
			p_pct = int(1 + (99 * ratio))

		try:
			with open(f"{PWM_PATH}/duty_cycle", "w") as f: f.write(str(duty))
		except: pass

		draw.text((2, 53), f"CPU:{int(t_c)}° DSK:{int(t_d)}° FAN:{p_pct}%", font=font, fill=1)
		send_fb(img)

	except Exception as e:
		print(f"[ERR]: {e}")
	time.sleep(1)