main.py

#!/usr/bin/env python3
"""
ZeroCD v2.0 — SuperDrive Changer HUD
"""
import signal
import sys
import os
import time
import subprocess

from config import ISO_DIR, BACKLIGHT_TIMEOUT_SECONDS, BACKLIGHT_FADE_STEPS, BACKLIGHT_FADE_STEP_MS, PRESET_IMG_SIZES
from system.logger import setup_logger, get_logger

from ui.display import Display
from input.joystick import Joystick, Direction
from usb.iso_manager import ISOManager
from net.wifi import WiFiManager
from usb.gadget import GadgetManager

MAX_SLOTS = 3


class ZeroCDApp:
    def __init__(self):
        self.logger = get_logger("main")
        self.display = None
        self.joystick = None
        self.gadget = None
        self.iso_manager = None
        self.wifi = None

        self.slot_paths = [None, None, None]
        self.slot_names = [None, None, None]
        self.active_slot_idx = -1
        self.cursor_slot = 0
        self.usb_connected = False
        self.running = False
        self.last_activity_time = None
        self.backlight_on = True
        self.in_submenu = False        # "slot_action" | "select_iso" | "create_img"
        self._sub_selected = 0
        self._iso_list_cache = []
        self._pending_slot = 0          # какой слот редактируем
        self.mtp_enabled = False

    def _scan_slots(self):
        slots_dict = self.iso_manager.list_changer_slots()
        for i in range(MAX_SLOTS):
            fname = slots_dict.get(i)
            if fname:
                path = self.iso_manager.get_iso_path(fname)
                self.slot_paths[i] = path
                self.slot_names[i] = fname
            else:
                self.slot_paths[i] = None
                self.slot_names[i] = None

    def init(self) -> bool:
        self.logger.info("ZeroCD v2.0 initializing")
        setup_logger()

        self.iso_manager = ISOManager(ISO_DIR)
        self._scan_slots()

        self.display = Display()
        if not self.display.init():
            self.logger.error("Failed to initialize display")
            return False

        self.display.show_splash()
        time.sleep(1.5)

        self.joystick = Joystick(disp=self.display, callback=self._on_joystick)

        hooked = 0
        for pin_attr, idx in [('GPIO_KEY1_PIN', 0), ('GPIO_KEY2_PIN', 1), ('GPIO_KEY3_PIN', 2)]:
            pin = getattr(self.display, pin_attr, None)
            if pin and hasattr(pin, 'when_pressed'):
                pin.when_pressed = (lambda i=idx: self.on_key_event(i))
                hooked += 1
        self.logger.info(f"Physical KEY1-3 hooked: {hooked}")

        self.gadget = GadgetManager()
        if not self.gadget.init():
            self.logger.warning("USB gadget not available")
        else:
            if self.gadget.bind():
                self.usb_connected = True
                self.logger.info("Gadget bound to UDC")

        self.last_activity_time = self.get_time()
        self.backlight_on = True
        self.display.fade_in(target_duty=50, steps=BACKLIGHT_FADE_STEPS, step_delay_ms=BACKLIGHT_FADE_STEP_MS)

        self._autoload_first_slot()

        try:
            import threading as _th
            from web.server import start_webui
            t = _th.Thread(target=start_webui, args=(self,), daemon=True)
            t.start()
            self.logger.info("WebUI thread started")
        except Exception as e:
            self.logger.error(f"WebUI start failed: {e}")

        self.wifi = WiFiManager()
        self.update_hud()
        self.logger.info("ZeroCD v2.0 ready")
        return True

    def _autoload_first_slot(self):
        for i in range(MAX_SLOTS):
            p = self.slot_paths[i]
            if p:
                self._activate_slot_internal(i, p)
                break

    def _activate_slot_internal(self, idx: int, iso_path: str):
        if idx == self.active_slot_idx:
            return
        self.logger.info(f"Loading slot {idx}: {os.path.basename(iso_path)}")
        if self.gadget:
            self.gadget.set_iso(iso_path)
            if self.usb_connected:
                self.gadget.switch_slot(idx)
        self.active_slot_idx = idx

    def on_key_event(self, slot_idx: int):
        self.logger.info(f"K{slot_idx+1} pressed")
        self.reset_activity()
        # If currently in submenu, use K1-3 as navigation shortcuts:
        #    K1 -> UP, K2 -> DOWN, K3 -> PRESS/select
        if self.in_submenu:
            if slot_idx == 0:
                self._on_joystick(Direction.UP)
            elif slot_idx == 1:
                self._on_joystick(Direction.DOWN)
            else:
                self._on_joystick(Direction.PRESS)
            return
        # Move cursor to the pressed slot and activate
        self.cursor_slot = slot_idx
        self._slot_action(slot_idx)

    def _slot_action(self, slot_idx: int):
        path = self.slot_paths[slot_idx]
        if path:
            self._activate_slot_internal(slot_idx, path)
            return

        self._pending_slot = slot_idx
        self.in_submenu = "slot_action"
        self._sub_selected = 0
        self.update_hud()

    # ---- SUBMENU SYSTEM (slot_action → select_iso | create_img → iso/files/sizes) ----

    def _exit_submenu(self):
        self.in_submenu = False
        self._pending_slot = 0
        self._sub_selected = 0
        self._iso_list_cache.clear()
        self.update_hud()

    def _refresh_iso_list(self):
        self._iso_list_cache.clear()
        d = os.path.join(ISO_DIR, '')
        # walk only files ending in .img or .iso
        for fname in sorted(os.listdir(ISO_DIR)):
            if fname.lower().endswith(('.img', '.iso')):
                self._iso_list_cache.append(fname)

    def _on_joystick(self, direction):
        self.reset_activity()
        self.logger.debug(f"JS direction={direction.value}")

        if self.in_submenu:
            self._handle_submenu(direction)
            return

        if direction == Direction.UP:
            self.cursor_slot = (self.cursor_slot - 1) % MAX_SLOTS
        elif direction == Direction.DOWN:
            self.cursor_slot = (self.cursor_slot + 1) % MAX_SLOTS
        elif direction == Direction.LEFT:
            self.toggle_mtp()
        elif direction == Direction.RIGHT:
            self.toggle_wifi()
        elif direction == Direction.PRESS:
            self._slot_action(self.cursor_slot)

        self.update_hud()

    def _handle_submenu(self, direction):
        if self.in_submenu == "slot_action":
            items = ["Select Image", "Create IMG"]
            if direction == Direction.UP:
                self._sub_selected = (self._sub_selected - 1) % 2
            elif direction == Direction.DOWN:
                self._sub_selected = (self._sub_selected + 1) % 2
            elif direction == Direction.LEFT:
                self._exit_submenu()
            elif direction == Direction.PRESS:
                if self._sub_selected == 0:
                    self._refresh_iso_list()
                    if not self._iso_list_cache:
                        self._exit_submenu()
                        return
                    self.in_submenu = "select_iso"
                    self._sub_selected = 0
                else:
                    self.in_submenu = "create_img"
                    self._sub_selected = 0

        elif self.in_submenu == "select_iso":
            items = self._iso_list_cache
            if not items:
                self._exit_submenu()
                return
            if direction == Direction.UP:
                self._sub_selected = (self._sub_selected - 1) % len(items)
            elif direction == Direction.DOWN:
                self._sub_selected = (self._sub_selected + 1) % len(items)
            elif direction == Direction.LEFT:
                self.in_submenu = "slot_action"
                self._sub_selected = 0
            elif direction == Direction.PRESS:
                fname = items[self._sub_selected]
                full = os.path.join(ISO_DIR, fname)
                if not os.path.isfile(full):
                    self.logger.error(f"File not found: {full}")
                    self._exit_submenu()
                    return
                # Попишем путь в очередь слота
                slot = self._pending_slot
                self.slot_paths[slot] = full
                self.slot_names[slot] = fname
                self._activate_slot_internal(slot, full)
                self._exit_submenu()

        elif self.in_submenu == "create_img":
            sizes = PRESET_IMG_SIZES
            if direction == Direction.LEFT:
                self.in_submenu = "slot_action"
                self._sub_selected = 0
            elif direction == Direction.UP:
                self._sub_selected = (self._sub_selected - 1) % len(sizes)
            elif direction == Direction.DOWN:
                self._sub_selected = (self._sub_selected + 1) % len(sizes)
            elif direction == Direction.PRESS:
                self._do_create_img()
                self._exit_submenu()

        self.update_hud()

    def _do_create_img(self):
        size_mb = PRESET_IMG_SIZES[self._sub_selected]['mb']
        name = self.iso_manager.get_next_disk_name()
        self.logger.info(f"Creating image: {name} ({size_mb}MB)")
        result = self.iso_manager.create_image(name, size_mb)
        if result:
            self.logger.info(f"IMG created: {result}")
            self._scan_slots()
        else:
            self.logger.error("Create image failed")

    # ---- WiFi & MTP --------------------------------------------------------

    def toggle_wifi(self):
        self.logger.info("WiFi toggle")
        if not self.wifi or not self.wifi.has_wifi_support():
            return
        try:
            from net.wifi import WiFiState
            s = self.wifi.get_status()
            if s == WiFiState.CONNECTED:
                self.wifi.disconnect()
                self.logger.info("WiFi disconnected")
            elif s in (WiFiState.OFF, WiFiState.ERROR):
                self.wifi.connect()
                self.logger.info("WiFi connect triggered")
            elif s == WiFiState.AP_MODE:
                self.wifi.stop_ap_mode()
                self.wifi.connect()
        except Exception as e:
            self.logger.error(f"WiFi toggle fail: {e}")

    def toggle_mtp(self):
        self.logger.info("MTP toggle requested")

        MTP_COMMAND = "umtprd"

        if not self.mtp_enabled:
            self.logger.info("Switching to MTP mode...")

            if self.gadget:
                self.gadget.shutdown()
                self.usb_connected = False

            time.sleep(0.5)

            self.logger.info("Setting up MTP USB Gadget...")
            os.system("sudo mkdir -p /sys/kernel/config/usb_gadget/mtp")
            os.system("sudo sh -c 'echo 0x1D6B > /sys/kernel/config/usb_gadget/mtp/idVendor'")
            os.system("sudo sh -c 'echo 0x0100 > /sys/kernel/config/usb_gadget/mtp/idProduct'")

            os.system("sudo mkdir -p /sys/kernel/config/usb_gadget/mtp/strings/0x409")
            os.system("sudo sh -c 'echo \"ZeroCD\" > /sys/kernel/config/usb_gadget/mtp/strings/0x409/manufacturer'")
            os.system("sudo sh -c 'echo \"ZeroCD MTP\" > /sys/kernel/config/usb_gadget/mtp/strings/0x409/product'")

            os.system("sudo mkdir -p /sys/kernel/config/usb_gadget/mtp/configs/c.1/strings/0x409")
            os.system("sudo sh -c 'echo \"MTP\" > /sys/kernel/config/usb_gadget/mtp/configs/c.1/strings/0x409/configuration'")

            os.system("sudo mkdir -p /sys/kernel/config/usb_gadget/mtp/functions/ffs.mtp")
            os.system("sudo ln -s /sys/kernel/config/usb_gadget/mtp/functions/ffs.mtp /sys/kernel/config/usb_gadget/mtp/configs/c.1/")

            os.system("sudo mkdir -p /dev/ffs-mtp")
            os.system("sudo mount -t functionfs mtp /dev/ffs-mtp")

            try:
                self.mtp_process = subprocess.Popen(["sudo", MTP_COMMAND])
                self.mtp_enabled = True
                self.logger.info("uMTP-Responder started. Waiting for endpoints...")

                time.sleep(1)

                udc_list = os.listdir("/sys/class/udc/")
                if udc_list:
                    udc = udc_list[0]
                    os.system(f"sudo sh -c 'echo {udc} > /sys/kernel/config/usb_gadget/mtp/UDC'")
                    self.logger.info(f"MTP connected to USB ({udc})")

            except Exception as e:
                self.logger.error(f"Failed to start MTP: {e}")
                self.mtp_enabled = False

        else:
            self.logger.info("Stopping MTP mode and restoring CD-ROM...")

            os.system("sudo sh -c 'echo \"\" > /sys/kernel/config/usb_gadget/mtp/UDC'")
            time.sleep(0.5)

            if hasattr(self, 'mtp_process') and self.mtp_process:
                self.mtp_process.terminate()
                os.system(f"sudo pkill -9 {MTP_COMMAND}")
                self.mtp_enabled = False

            os.system("sudo umount /dev/ffs-mtp")
            os.system("sudo rm /sys/kernel/config/usb_gadget/mtp/configs/c.1/ffs.mtp")
            os.system("sudo rmdir /sys/kernel/config/usb_gadget/mtp/configs/c.1/strings/0x409")
            os.system("sudo rmdir /sys/kernel/config/usb_gadget/mtp/configs/c.1")
            os.system("sudo rmdir /sys/kernel/config/usb_gadget/mtp/functions/ffs.mtp")
            os.system("sudo rmdir /sys/kernel/config/usb_gadget/mtp/strings/0x409")
            os.system("sudo rmdir /sys/kernel/config/usb_gadget/mtp")

            time.sleep(0.5)

            if self.gadget:
                if self.gadget.init():
                    if self.gadget.bind():
                        self.usb_connected = True

            self.logger.info("Refreshing ISO list after MTP session...")
            self._scan_slots()
            self.update_hud()

    # ---- DISPLAY HUD -------------------------------------------------------

    def update_hud(self):
        if not self.display:
            return

        if self.in_submenu == "slot_action":
            self.display.draw_submenu(
                title="Slot action",
                items=["Select Image", "Create IMG"],
                selected_idx=self._sub_selected,
                footer="LEFT=back PRESS=select"
            )
            return

        if self.in_submenu == "select_iso":
            self.display.draw_submenu(
                title="Select Image",
                items=self._iso_list_cache,
                selected_idx=self._sub_selected,
                footer="LEFT=back PRESS=select"
            )
            return

        if self.in_submenu == "create_img":
            self.display.draw_create_menu(
                [s['label'] for s in PRESET_IMG_SIZES],
                self._sub_selected
            )
            return

        self.display.set_slots(self.slot_names[:MAX_SLOTS])
        self.display.set_cursor(self.cursor_slot)
        self.display.set_active_slot(self.active_slot_idx)
        if self.active_slot_idx >= 0:
            self.display.set_scsi_status(f"SLOT {self.active_slot_idx+1}")
        else:
            self.display.set_scsi_status("NO DISC")
        self.display.draw_changer_hud()
        self.display.update()

    def reset_activity(self):
        self.last_activity_time = self.get_time()
        if not self.backlight_on:
            self.backlight_on = True
            self.display.fade_in(target_duty=50, steps=BACKLIGHT_FADE_STEPS, step_delay_ms=BACKLIGHT_FADE_STEP_MS)

    def get_time(self):
        return __import__('time').time()

    def check_backlight_timeout(self):
        if not self.backlight_on:
            return
        if self.get_time() - self.last_activity_time >= BACKLIGHT_TIMEOUT_SECONDS:
            self.backlight_on = False
            self.display.fade_out(steps=BACKLIGHT_FADE_STEPS, step_delay_ms=BACKLIGHT_FADE_STEP_MS)

    def run(self):
        self.running = True
        try:
            self.joystick.start_polling(self._on_joystick)
        except Exception as e:
            self.logger.warning(f"Joystick unavailable: {e}")
        import time as _t
        while self.running:
            self.check_backlight_timeout()
            _t.sleep(0.5)

    def shutdown(self):
        self.running = False
        if self.joystick:
            self.joystick.stop()
        if self.gadget:
            self.gadget.shutdown()
        if self.display:
            self.display.close()
            self.display = None


def main():
    def signal_handler(signum, frame):
        global app
        if app:
            app.logger.info(f"Signal {signum}")
            app.shutdown()
        sys.exit(0)

    signal.signal(signal.SIGTERM, signal_handler)
    signal.signal(signal.SIGINT, signal_handler)

    global app
    app = ZeroCDApp()
    if app.init():
        app.run()
    else:
        sys.exit(1)


if __name__ == "__main__":
    main()