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


BASE_DIR = os.path.dirname(os.path.abspath(__file__))

PORT_MAP = {3: 0, 4: 1, 5: 2, 6: 3, 2: 4, 1: 5}  # ASM1166 port map
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 = {}
fb_file = open(FB_PATH, "r+b")
fb_map = mmap.mmap(fb_file.fileno(), W * H * 4)

try:
	font = ImageFont.truetype("/usr/share/fonts/X11/misc/5x7.pcf.gz", 7)
except Exception as e:
	print("[ERROR]", e)
	font = ImageFont.load_default()


# --- Вспомогательные функции ---
def load_res(name):
	try:
		p = os.path.join(BASE_DIR, "images", name)
		img = Image.open(p).convert('1')
		return ImageOps.invert(img)
	except Exception as e:
		print("[ERROR]", e)
		exit(1)


def send_fb(img):
	try:
		mask = np.array(img) > 0
		data = np.zeros((H, W), dtype=np.uint32)
		data[mask] = 0xFFFFFF00
		fb_map.seek(0)
		fb_map.write(data.tobytes())
	except Exception as e:
		print("[ERROR]", e)
		exit(2)


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 Exception as e:
		print("[ERROR]", e)
		exit(3)


# --- Thread 1: Disks activity ---
def disk_scanner_worker():
	global shared_slots
	last_io = {}

	while True:
		new_slots = [None] * 6
		for dev_path in glob.glob("/sys/block/sd*"):
			dev = os.path.basename(dev_path)

			if not os.path.exists(f"/dev/{dev}"):
				continue

			try:
				# /sys/block/sda -> ../devices/platform/soc@3000000/6000000.pcie/pci0000:00/0000:00:00.0/0000:01:00.0/ata3/host3/target3:0:0/3:0:0:0/block/sda
				real_p = os.readlink(dev_path)
				if "pcie" not in real_p:
					continue  # Skip USB disks (and other)

				# Ищем номер порта
				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:
						# /sys/block/sda/stat
						# 220 0 12880 517 0 0 0 0 0 504 517 0 0 0 0 0 0
						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 Exception as e:
				print("[ERROR]", e)
				continue

		# Disks temp (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 Exception as e:
				print("[ERROR]", e)

		# CPU temp
		c_t = 0
		for p in glob.glob("/sys/class/thermal/thermal_zone*/"):
			with open(p + "temp") as f:
				c_t = max(c_t, int(f.read()) / 1000)

		with lock:
			shared_slots = new_slots
			shared_temps['disk'] = d_t
			shared_temps['cpu'] = c_t

		time.sleep(0.5)


# --- Поток 2: 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:
				res = subprocess.run(['smartctl', '-H', f'/dev/{d}'], capture_output=True, text=True, timeout=5)
				# smartctl -H /dev/sda
				# === START OF READ SMART DATA SECTION ===
				# SMART overall-health self-assessment test result: PASSED
				shared_health[d] = "PASSED" in res.stdout
			except Exception as e:
				print("[ERROR]", e)
				pass
		time.sleep(60)


# Prepare resources
setup_pwm()

IMG_BOOT = load_res("boot.png")
IMG_SHUTDOWN = load_res("shutdown.png")
IMG_DISK_OK = load_res("disk_ok.png")
IMG_DISK_NO = load_res("disk_empty.png")
IMG_DISK_ERR = load_res("disk_error.png")
IMG_DISK_ACT = load_res("disk_active.png")
ICON_DISK = load_res("icon_disk.png")
ICON_CPU = load_res("icon_cpu.png")

signal.signal(signal.SIGTERM, lambda s, f: (send_fb(IMG_SHUTDOWN), sys.exit(0)))

# Boot animation
i = 0
while True:
	tmp = IMG_BOOT.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.25)
	try:
		status = subprocess.run(["systemctl", "is-system-running"], capture_output=True, text=True)
		if status.stdout.strip() in ("running", "degraded"):
			break
	except Exception as e:
		print("ERROR:", e)
		exit(4)


# Main thread
print("System ready, starting all threads...")
threading.Thread(target=disk_scanner_worker, daemon=True).start()
threading.Thread(target=smart_worker, daemon=True).start()
blink = True
fan_on = False
while True:
	try:
		img = Image.new('1', (W, H), 0)
		draw = ImageDraw.Draw(img)
		blink = not blink

		with lock:
			current_view = list(shared_slots)
			t_d = shared_temps['disk']
			t_c = shared_temps['cpu']

		# Draw time
		draw.text((0, 0), time.strftime("%H:%M:%S"), font=font, fill=1)

		# Draw CPU temp
		img.paste(ICON_CPU, (80, 0))
		draw.text((86, 0), f"{t_c:2.0f}°", font=font, fill=1)

		# Draw HDD temp
		img.paste(ICON_DISK, (104, 0))
		draw.text((110, 0), f"{t_d:2.0f}°", font=font, fill=1)

		# Draw separator/line
		draw.line((0, 8, 128, 8), fill=1)

		# Draw disks
		for i in range(6):
			slot = current_view[i]
			if slot:
				if not slot['ok']:
					ic = IMG_DISK_ERR if blink else IMG_DISK_NO
					# TODO: Turn ON backlight
				elif slot['act']:
					ic = IMG_DISK_ACT
				else:
					ic = IMG_DISK_OK
			else:
				ic = IMG_DISK_NO
			img.paste(ic, (2 + (i * 21), 10))

		# 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 = 33 + (i * 8)
			draw.text((2, y), l, font=font, fill=1)
			draw.rounded_rectangle((26, y, 102, y + 6), radius=1, outline=1)
			draw.rectangle((28, y + 2, 28 + int(72 * v / 100), y + 4), fill=1)
			draw.text((104, y), f"{int(v):3d}%", font=font, fill=1)
		send_fb(img)

		# Fan control
		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 Exception as e:
			print("[ERROR]", e)

	except Exception as e:
		print("[ERROR]", e)
	time.sleep(1)