"""
Loki — HID Attack Engine for Bjorn.

Manages USB HID gadget lifecycle, script execution, and job tracking.
Named after the Norse trickster god.

Loki is the 5th exclusive operation mode (alongside MANUAL, AUTO, AI, BIFROST).
When active, the orchestrator stops and the Pi acts as a keyboard/mouse
to the connected host via /dev/hidg0 (keyboard) and /dev/hidg1 (mouse).

HID GADGET STRATEGY:
  The HID functions (keyboard + mouse) are created ONCE at boot time alongside
  RNDIS networking by the usb-gadget.sh script. This avoids the impossible task
  of hot-adding HID functions to a running composite gadget (UDC rebind fails
  with EIO when RNDIS is active).

  LokiEngine simply opens/closes the /dev/hidg0 and /dev/hidg1 device files.
  If /dev/hidg0 doesn't exist, the user needs to run the setup once and reboot.
"""
import os
import time
import subprocess
import logging
from threading import Event

from logger import Logger

logger = Logger(name="loki", level=logging.DEBUG)

# USB HID report descriptors — EXACT byte-for-byte copies from P4wnP1_aloa
# Source: P4wnP1_aloa-master/service/SubSysUSB.go lines 54-70
#
# These are written to the gadget at boot time by usb-gadget.sh.
# Kept here for reference and for the install_hid_gadget() method.
#
_KBD_REPORT_DESC = bytes([
    0x05, 0x01, 0x09, 0x06, 0xa1, 0x01, 0x05, 0x07,
    0x19, 0xe0, 0x29, 0xe7, 0x15, 0x00, 0x25, 0x01,
    0x75, 0x01, 0x95, 0x08, 0x81, 0x02, 0x95, 0x01,
    0x75, 0x08, 0x81, 0x03, 0x95, 0x05, 0x75, 0x01,
    0x05, 0x08, 0x19, 0x01, 0x29, 0x05, 0x91, 0x02,
    0x95, 0x01, 0x75, 0x03, 0x91, 0x03, 0x95, 0x06,
    0x75, 0x08, 0x15, 0x00, 0x25, 0x65, 0x05, 0x07,
    0x19, 0x00, 0x29, 0x65, 0x81, 0x00, 0xc0,
])  # 63 bytes, report_length=8

_MOUSE_REPORT_DESC = bytes([
    0x05, 0x01, 0x09, 0x02, 0xa1, 0x01, 0x09, 0x01,
    0xa1, 0x00, 0x85, 0x01, 0x05, 0x09, 0x19, 0x01,
    0x29, 0x03, 0x15, 0x00, 0x25, 0x01, 0x95, 0x03,
    0x75, 0x01, 0x81, 0x02, 0x95, 0x01, 0x75, 0x05,
    0x81, 0x03, 0x05, 0x01, 0x09, 0x30, 0x09, 0x31,
    0x15, 0x81, 0x25, 0x7f, 0x75, 0x08, 0x95, 0x02,
    0x81, 0x06, 0x95, 0x02, 0x75, 0x08, 0x81, 0x01,
    0xc0, 0xc0, 0x05, 0x01, 0x09, 0x02, 0xa1, 0x01,
    0x09, 0x01, 0xa1, 0x00, 0x85, 0x02, 0x05, 0x09,
    0x19, 0x01, 0x29, 0x03, 0x15, 0x00, 0x25, 0x01,
    0x95, 0x03, 0x75, 0x01, 0x81, 0x02, 0x95, 0x01,
    0x75, 0x05, 0x81, 0x01, 0x05, 0x01, 0x09, 0x30,
    0x09, 0x31, 0x15, 0x00, 0x26, 0xff, 0x7f, 0x95,
    0x02, 0x75, 0x10, 0x81, 0x02, 0xc0, 0xc0,
])  # 111 bytes, report_length=6

# The boot script that creates RNDIS + HID functions at startup.
# This replaces /usr/local/bin/usb-gadget.sh
_USB_GADGET_SCRIPT = '''#!/bin/bash
# usb-gadget.sh — USB composite gadget: RNDIS networking + HID (keyboard/mouse)
# Auto-generated by Bjorn Loki. Do not edit manually.

modprobe libcomposite
cd /sys/kernel/config/usb_gadget/
mkdir -p g1
cd g1

echo 0x1d6b > idVendor
echo 0x0104 > idProduct
echo 0x0100 > bcdDevice
echo 0x0200 > bcdUSB

mkdir -p strings/0x409
echo "fedcba9876543210" > strings/0x409/serialnumber
echo "Raspberry Pi" > strings/0x409/manufacturer
echo "Pi Zero USB" > strings/0x409/product

mkdir -p configs/c.1/strings/0x409
echo "Config 1: RNDIS + HID" > configs/c.1/strings/0x409/configuration
echo 250 > configs/c.1/MaxPower

# ── RNDIS networking function ──
mkdir -p functions/rndis.usb0
[ -L configs/c.1/rndis.usb0 ] && rm configs/c.1/rndis.usb0
ln -s functions/rndis.usb0 configs/c.1/

# ── HID functions (keyboard + mouse) ──
# Use python3 to write binary report descriptors (bash printf can't handle null bytes)
python3 - <<'PYEOF'
import os, sys

G = "/sys/kernel/config/usb_gadget/g1"

# Keyboard: P4wnP1 exact boot keyboard descriptor (63 bytes)
KBD_DESC = bytes([
    0x05,0x01,0x09,0x06,0xa1,0x01,0x05,0x07,
    0x19,0xe0,0x29,0xe7,0x15,0x00,0x25,0x01,
    0x75,0x01,0x95,0x08,0x81,0x02,0x95,0x01,
    0x75,0x08,0x81,0x03,0x95,0x05,0x75,0x01,
    0x05,0x08,0x19,0x01,0x29,0x05,0x91,0x02,
    0x95,0x01,0x75,0x03,0x91,0x03,0x95,0x06,
    0x75,0x08,0x15,0x00,0x25,0x65,0x05,0x07,
    0x19,0x00,0x29,0x65,0x81,0x00,0xc0,
])

# Mouse: P4wnP1 dual-mode (relative + absolute) descriptor (111 bytes)
MOUSE_DESC = bytes([
    0x05,0x01,0x09,0x02,0xa1,0x01,0x09,0x01,
    0xa1,0x00,0x85,0x01,0x05,0x09,0x19,0x01,
    0x29,0x03,0x15,0x00,0x25,0x01,0x95,0x03,
    0x75,0x01,0x81,0x02,0x95,0x01,0x75,0x05,
    0x81,0x03,0x05,0x01,0x09,0x30,0x09,0x31,
    0x15,0x81,0x25,0x7f,0x75,0x08,0x95,0x02,
    0x81,0x06,0x95,0x02,0x75,0x08,0x81,0x01,
    0xc0,0xc0,0x05,0x01,0x09,0x02,0xa1,0x01,
    0x09,0x01,0xa1,0x00,0x85,0x02,0x05,0x09,
    0x19,0x01,0x29,0x03,0x15,0x00,0x25,0x01,
    0x95,0x03,0x75,0x01,0x81,0x02,0x95,0x01,
    0x75,0x05,0x81,0x01,0x05,0x01,0x09,0x30,
    0x09,0x31,0x15,0x00,0x26,0xff,0x7f,0x95,
    0x02,0x75,0x10,0x81,0x02,0xc0,0xc0,
])

def w(path, content):
    with open(path, "w") as f:
        f.write(content)

def wb(path, data):
    with open(path, "wb") as f:
        f.write(data)

try:
    # Keyboard (hid.usb0)
    kbd = G + "/functions/hid.usb0"
    os.makedirs(kbd, exist_ok=True)
    w(kbd + "/protocol", "1")
    w(kbd + "/subclass", "1")
    w(kbd + "/report_length", "8")
    wb(kbd + "/report_desc", KBD_DESC)

    # Mouse (hid.usb1)
    mouse = G + "/functions/hid.usb1"
    os.makedirs(mouse, exist_ok=True)
    w(mouse + "/protocol", "2")
    w(mouse + "/subclass", "1")
    w(mouse + "/report_length", "6")
    wb(mouse + "/report_desc", MOUSE_DESC)

    print(f"HID functions created: keyboard ({len(KBD_DESC)}B) + mouse ({len(MOUSE_DESC)}B)")
except Exception as e:
    print(f"WARNING: HID setup failed (non-fatal): {e}", file=sys.stderr)
    sys.exit(0)  # Don't block RNDIS setup
PYEOF

# Symlink HID into config (non-fatal)
for func in hid.usb0 hid.usb1; do
    [ -L "configs/c.1/$func" ] && rm "configs/c.1/$func"
    if [ -d "functions/$func" ]; then
        ln -s "functions/$func" "configs/c.1/" 2>/dev/null || true
    fi
done

# ── Bind UDC ──
max_retries=10
retry_count=0
while ! ls /sys/class/udc > UDC 2>/dev/null; do
    if [ $retry_count -ge $max_retries ]; then
        echo "Error: Device or resource busy after $max_retries attempts."
        exit 1
    fi
    retry_count=$((retry_count + 1))
    sleep 1
done

UDC_NAME=$(ls /sys/class/udc)
echo "$UDC_NAME" > UDC
echo "Assigned UDC: $UDC_NAME (RNDIS + HID keyboard + HID mouse)"

# ── Configure network ──
if ! ip addr show usb0 2>/dev/null | grep -q "172.20.2.1"; then
    ifconfig usb0 172.20.2.1 netmask 255.255.255.0 2>/dev/null || true
    echo "Configured usb0 with IP 172.20.2.1"
else
    echo "Interface usb0 already configured."
fi
'''

_GADGET_SCRIPT_PATH = "/usr/local/bin/usb-gadget.sh"


class LokiEngine:
    """HID attack engine — manages script execution and job tracking.

    The USB HID gadget (keyboard + mouse) is set up at boot time by
    usb-gadget.sh. This engine simply opens /dev/hidg0 and /dev/hidg1.
    """

    def __init__(self, shared_data):
        self.shared_data = shared_data
        self._stop_event = Event()
        self._running = False
        self._gadget_ready = False

        # Sub-components (lazy init)
        self._hid = None
        self._jobs = None

    # ── Properties ─────────────────────────────────────────────

    @property
    def hid_controller(self):
        if self._hid is None:
            from loki.hid_controller import HIDController
            self._hid = HIDController()
        return self._hid

    @property
    def job_manager(self):
        if self._jobs is None:
            from loki.jobs import LokiJobManager
            self._jobs = LokiJobManager(self)
        return self._jobs

    # ── Start / Stop ───────────────────────────────────────────

    def start(self):
        """Start Loki engine: open HID devices, apply config."""
        if self._running:
            logger.warning("Loki already running")
            return

        logger.info("Starting Loki engine...")
        self._stop_event.clear()

        # Check if HID gadget is available (set up at boot)
        if not os.path.exists("/dev/hidg0"):
            logger.error(
                "/dev/hidg0 not found — HID gadget not configured at boot. "
                "Run install_hid_gadget() from the Loki API and reboot."
            )
            self._gadget_ready = False
            return

        self._gadget_ready = True

        # Open HID devices
        try:
            self.hid_controller.open()
        except Exception as e:
            logger.error("HID device open failed: %s", e)
            self._gadget_ready = False
            return

        # Apply config
        layout = self.shared_data.config.get("loki_default_layout", "us")
        self.hid_controller.set_layout(layout)
        speed_min = self.shared_data.config.get("loki_typing_speed_min", 0)
        speed_max = self.shared_data.config.get("loki_typing_speed_max", 0)
        self.hid_controller.set_typing_speed(speed_min, speed_max)

        self._running = True
        logger.info("Loki engine started (HID devices open)")

        # Auto-run script if configured
        auto_run = self.shared_data.config.get("loki_auto_run", "")
        if auto_run:
            self._auto_run_script(auto_run)

    def stop(self):
        """Stop Loki engine: cancel jobs, close devices."""
        if not self._running:
            return

        logger.info("Stopping Loki engine...")
        self._stop_event.set()

        # Cancel all running jobs
        if self._jobs:
            for job in self._jobs.get_all_jobs():
                if job["status"] == "running":
                    self._jobs.cancel_job(job["id"])

        # Close HID devices (don't remove gadget — it persists)
        if self._hid:
            self._hid.close()

        self._running = False
        self._gadget_ready = False
        logger.info("Loki engine stopped")

    def get_status(self) -> dict:
        """Return current engine status for the API."""
        hidg0_exists = os.path.exists("/dev/hidg0")
        return {
            "enabled": self.shared_data.config.get("loki_enabled", False),
            "running": self._running,
            "gadget_ready": self._gadget_ready,
            "gadget_installed": hidg0_exists,
            "layout": self.shared_data.config.get("loki_default_layout", "us"),
            "jobs_running": self.job_manager.running_count if self._jobs else 0,
            "jobs_total": len(self._jobs.get_all_jobs()) if self._jobs else 0,
        }

    # ── Job API (delegated to JobManager) ──────────────────────

    def submit_job(self, script_name: str, script_content: str) -> str:
        """Submit a HIDScript for execution. Returns job_id."""
        if not self._running:
            raise RuntimeError("Loki engine not running")
        if not self._gadget_ready:
            raise RuntimeError("HID gadget not ready")
        return self.job_manager.create_job(script_name, script_content)

    def cancel_job(self, job_id: str) -> bool:
        return self.job_manager.cancel_job(job_id)

    def get_jobs(self) -> list:
        return self.job_manager.get_all_jobs()

    # ── HID Gadget Installation ────────────────────────────────

    @staticmethod
    def is_gadget_installed() -> bool:
        """Check if the HID gadget is available."""
        return os.path.exists("/dev/hidg0")

    @staticmethod
    def install_hid_gadget() -> dict:
        """Install/update the USB gadget boot script to include HID functions.

        Writes the new usb-gadget.sh that creates RNDIS + HID at boot.
        Returns status dict. Requires a reboot to take effect.
        """
        try:
            # Write the new gadget script
            with open(_GADGET_SCRIPT_PATH, "w") as f:
                f.write(_USB_GADGET_SCRIPT)
            os.chmod(_GADGET_SCRIPT_PATH, 0o755)
            logger.info("USB gadget script updated at %s", _GADGET_SCRIPT_PATH)

            return {
                "success": True,
                "message": "USB gadget script updated with HID support. Reboot required.",
                "reboot_required": True,
            }

        except Exception as e:
            logger.error("Failed to install HID gadget script: %s", e)
            return {
                "success": False,
                "message": f"Installation failed: {e}",
                "reboot_required": False,
            }

    # ── Auto-run ───────────────────────────────────────────────

    def _auto_run_script(self, script_name: str):
        """Auto-run a script by name from the database."""
        try:
            db = self.shared_data.db
            row = db.query_one(
                "SELECT content FROM loki_scripts WHERE name = ?", (script_name,)
            )
            if row and row.get("content"):
                self.submit_job(script_name, row["content"])
                logger.info("Auto-running script: %s", script_name)
        except Exception as e:
            logger.error("Auto-run failed for '%s': %s", script_name, e)