import time
import os
import psutil
import socket
import subprocess
from datetime import datetime
from collections import deque
from luma.core.interface.serial import i2c
from luma.oled.device import ssd1306
from PIL import Image, ImageDraw, ImageFont

# ---------------- CONSTANTS ----------------
I2C_PORT = 3
I2C_ADDR = 0x3C

CPU_HISTORY_LEN = 100
UPDATE_INTERVAL = 0.5

GRAPH_HEIGHT = 30
GRAPH_FOOTER_HEIGHT = 10

RAM_BAR_WIDTH = 25
RAM_BAR_HEIGHT = 52

FS_BOX = (64, 42, 100, 63)
FS_PIE_CENTER = (74, 52)
FS_PIE_RADIUS = 9

scroll_offset = 0
SCROLL_SPEED = 2  # pixels per update

# ---------------- OLED SETUP ----------------
serial = i2c(port=I2C_PORT, address=I2C_ADDR)
device = ssd1306(serial)
WIDTH, HEIGHT = device.width, device.height

cpu_history = deque([0] * CPU_HISTORY_LEN, maxlen=CPU_HISTORY_LEN)

try:
    small_font = ImageFont.truetype("DejaVuSans.ttf", 10)
except IOError:
    small_font = ImageFont.load_default()

# ---------------- HELPERS ----------------
def get_cpu_temp():
    temps = psutil.sensors_temperatures()
    if temps:
        return next(iter(temps.values()))[0].current
    try:
        with open("/sys/class/thermal/thermal_zone0/temp") as f:
            return int(f.read()) / 1000
    except OSError:
        return 0.0

def get_root_fs_usage():
    return psutil.disk_usage("/").percent

def format_uptime():
    boot = datetime.fromtimestamp(psutil.boot_time())
    seconds = int((datetime.now() - boot).total_seconds())
    hours, rem = divmod(seconds, 3600)
    minutes, _ = divmod(rem, 60)
    return hours, minutes

def truncate_text(text, max_len=8):
    if len(text) <= max_len:
        return text
    return text[:max_len - 1] + "…"

def get_wifi_info(interface="wlan0"):
    try:
        result = subprocess.run(
            ["iw", "dev", interface, "link"],
            capture_output=True,
            text=True
        )

        if "Not connected" in result.stdout:
            return "--", 0

        ssid = "--"
        signal_dbm = None

        for line in result.stdout.splitlines():
            line = line.strip()

            if line.startswith("SSID:"):
                ssid = line.split("SSID:")[1].strip()

            if line.startswith("signal:"):
                signal_dbm = int(line.split()[1])

        # Convert dBm → percentage (rough but standard)
        if signal_dbm is not None:
            signal_pct = max(0, min(100, 2 * (signal_dbm + 100)))
        else:
            signal_pct = 0

        return truncate_text(ssid), signal_pct

    except Exception:
        return "--", 0

def estimate_power_w(cpu_percent):
    # Constants (from measurement)
    BASE_POWER = 1.9          # W
    IDLE_FREQ_POWER = 0.3     # W at max freq
    MAX_CPU_POWER = 3.0       # W at full load + max freq

    MIN_FREQ = 614.0          # MHz
    MAX_FREQ = 1600.0         # MHz

    # Read current frequency (assume all cores similar)
    with open("/sys/devices/system/cpu/cpu0/cpufreq/scaling_cur_freq") as f:
        freq_mhz = int(f.read()) / 1000

    freq_factor = max(
        0.0,
        min(1.0, (freq_mhz - MIN_FREQ) / (MAX_FREQ - MIN_FREQ))
    )

    load = cpu_percent / 100.0

    power = (
        BASE_POWER +
        IDLE_FREQ_POWER * freq_factor +
        MAX_CPU_POWER * freq_factor * load
    )

    return round(power, 2)

# ---------------- DRAWING ----------------
def draw_cpu_graph(draw, cpu, load, temp, watts):
    draw.rectangle((0, 0, CPU_HISTORY_LEN, GRAPH_HEIGHT), outline=255)

    if len(cpu_history) > 1:
        points = [
            (i, GRAPH_HEIGHT - (v / 100 * GRAPH_HEIGHT))
            for i, v in enumerate(cpu_history)
        ]
        points += [(len(cpu_history) - 1, GRAPH_HEIGHT), (0, GRAPH_HEIGHT)]
        draw.polygon(points, fill=255)

    footer_top = GRAPH_HEIGHT
    footer_bottom = GRAPH_HEIGHT + GRAPH_FOOTER_HEIGHT

    draw.rectangle((0, footer_top, CPU_HISTORY_LEN, footer_bottom), outline=255)
    draw.rectangle((0, footer_top, 25, footer_bottom), fill=255)

    draw.text((1, footer_top - 1), "CPU", font=small_font, fill=0)
    draw.line((59, footer_top, 59, footer_bottom), fill=255, width=2)

    load_text = f"{watts:.1f} W"
    lw = draw.textbbox((0, 0), load_text, font=small_font)[2]
    load_x = 25 + (59 - 25 - lw) // 2
    draw.text((load_x, footer_top - 1), load_text, font=small_font, fill=255)

    cpu_text = f"{cpu:.0f}%"
    cw = draw.textbbox((0, 0), cpu_text, font=small_font)[2]
    cpu_x = 59 + (CPU_HISTORY_LEN - 59 - cw) // 2
    draw.text((cpu_x, footer_top - 1), cpu_text, font=small_font, fill=255)

    draw.rectangle((0, 0, 26, 10), fill=255)
    draw.text((1, -1), f"{temp:.0f}°C", font=small_font, fill=0)

def draw_ram_bar(draw, mem_percent):
    x = WIDTH - 1 - RAM_BAR_WIDTH
    y = 0

    draw.rectangle((x, y, WIDTH - 1, RAM_BAR_HEIGHT), outline=255)

    fill_h = int((mem_percent / 100) * RAM_BAR_HEIGHT)
    draw.rectangle(
        (x, RAM_BAR_HEIGHT - fill_h, WIDTH - 1, RAM_BAR_HEIGHT),
        fill=255
    )

    label = f"{mem_percent:.0f}%"
    tw = draw.textbbox((0, 0), label, font=small_font)[2]
    draw.text((x + (RAM_BAR_WIDTH - tw) // 2 + 1,
               RAM_BAR_HEIGHT - fill_h - 11),
              label, font=small_font, fill=255)

    draw.rectangle((x, RAM_BAR_HEIGHT, WIDTH - 1, HEIGHT - 1), outline=0, fill=255)
    draw.text((103, 52), "RAM", font=small_font, fill=0)

def draw_fs_pie(draw, used_pct):
    draw.rectangle(FS_BOX, outline=255)

    cx, cy = FS_PIE_CENTER
    r = FS_PIE_RADIUS
    bbox = (cx - r, cy - r, cx + r, cy + r)

    draw.ellipse(bbox, outline=255)
    angle = int(360 * (used_pct / 100))

    draw.pieslice(bbox, start=270, end=270 + angle, fill=255)
    draw.text((86, 42), "FS", font=small_font, fill=255)
    draw.text((90, 51), "/", font=small_font, fill=255)

def draw_uptime(draw, hours, minutes):
    draw.rectangle((-1, 53, 62, HEIGHT), outline=255)
    draw.text((0, 53), f"up: {hours}h {minutes}m",
              font=small_font, fill=255)

def draw_wifi_icon(draw, ssid, signal_pct, x=0, y=0):
    bars = 4
    bar_width = 2
    bar_spacing = 1
    max_height = 8   # tallest bar
    filled_bars = int(signal_pct / 100 * bars + 0.5)

    for i in range(bars):
        bar_height = int((i + 1) / bars * max_height)
        y_top = y + (max_height - bar_height)
        fill = 1 if i < filled_bars else 0
        draw.rectangle(
            (x + i * (bar_width + bar_spacing), y_top,
             x + i * (bar_width + bar_spacing) + bar_width,
             y + max_height),
            fill=fill
        )

    # Draw SSID text right next to icon
    text_x = x + bars * (bar_width + bar_spacing) + 1
    draw.text((text_x, y-3), ssid, font=small_font, fill=255)

# ---------------- MAIN UPDATE ----------------
def update_display():
    cpu = psutil.cpu_percent(interval=0.1)
    watts = estimate_power_w(cpu)
    mem = psutil.virtual_memory().percent
    load1, _, _ = os.getloadavg()
    cpu_history.append(cpu)

    temp = get_cpu_temp()
    fs_used = get_root_fs_usage()
    hours, minutes = format_uptime()

    image = Image.new("1", (WIDTH, HEIGHT))
    draw = ImageDraw.Draw(image)

    draw.rectangle((0, 0, WIDTH, HEIGHT), fill=0)

    ssid, signal = get_wifi_info("wlan0")
    draw_wifi_icon(draw, ssid, signal, x=0, y=43)
    draw_cpu_graph(draw, cpu, load1, temp, watts)
    draw_ram_bar(draw, mem)
    draw_fs_pie(draw, fs_used)
    draw_uptime(draw, hours, minutes)

    device.display(image)

# ---------------- MAIN LOOP ----------------
if __name__ == "__main__":
    try:
        for _ in range(CPU_HISTORY_LEN):
            cpu_history.append(psutil.cpu_percent(interval=0.01))

        while True:
            update_display()
            time.sleep(UPDATE_INTERVAL)
    except KeyboardInterrupt:
        device.clear()