#!/usr/bin/env python3
# ============================================================================
#  _gkfit.py  —  shared motion-input helper for GAME KING FIT titles
#
#  Games launched from the Fit menu are separate processes, so they don't
#  inherit the menu's pointer. Import this instead and you get the same
#  absolute gyro aiming, the same buttons, and the same settings file:
#
#      from _gkfit import GKInput
#      gk = GKInput(W, H)
#      ...
#      gk.update(dt)
#      x, y = gk.pos
#      if "a" in gk.pressed: shoot()
#      if gk.swing > 200: bowl(gk.swing)
#
#  Falls back to mouse + keyboard + gamepad when there's no Joy-Con, so every
#  title stays playable without one. Leading underscore keeps the launcher
#  from listing this as a game.
# ============================================================================

import os
import json
import math
import time

import pygame

try:
    import evdev
    from evdev import ecodes
    HAVE_EVDEV = True
except Exception:
    HAVE_EVDEV = False

CONF_PATH = os.path.expanduser("~/.gameking_fit.json")

GYRO_MAPS = [("z", "y"), ("z", "x"), ("x", "y"), ("x", "z"),
             ("y", "z"), ("y", "x")]

BTN_MAP = {}
if HAVE_EVDEV:
    BTN_MAP = {
        ecodes.BTN_EAST: "a", ecodes.BTN_SOUTH: "b",
        ecodes.BTN_NORTH: "x", ecodes.BTN_WEST: "y",
        ecodes.BTN_MODE: "home", ecodes.BTN_START: "plus",
        ecodes.BTN_SELECT: "minus",
        ecodes.BTN_TR: "r", ecodes.BTN_TR2: "zr",
        ecodes.BTN_TL: "r", ecodes.BTN_TL2: "zr",
        ecodes.BTN_DPAD_DOWN: "a", ecodes.BTN_DPAD_LEFT: "b",
        ecodes.BTN_DPAD_RIGHT: "x", ecodes.BTN_DPAD_UP: "y",
    }


def load_conf():
    c = {"sens": 1.0, "invert_x": False, "invert_y": False,
         "swap_axes": False, "gyro_map": 0, "pad_layout": "nintendo"}
    try:
        with open(CONF_PATH) as f:
            c.update(json.load(f))
    except Exception:
        pass
    return c


def open_screen(w, h, caption="GAME KING FIT", force_windowed=False):
    """Open a display that matches however the Fit console itself is running.

    The launcher exports GKFIT_FULLSCREEN, so a game started from a fullscreen
    console goes fullscreen and one started from a window stays windowed.
    Run standalone (no launcher), it defaults to a window. pygame.SCALED lets
    the game keep its own fixed logical size and be stretched to whatever the
    real panel is, so titles don't need to know the display resolution.
    """
    want_fs = os.environ.get("GKFIT_FULLSCREEN", "0") == "1"
    if force_windowed or os.environ.get("GKFIT_WINDOWED") == "1":
        want_fs = False
    flags = pygame.SCALED
    if want_fs:
        flags |= pygame.FULLSCREEN
    try:
        screen = pygame.display.set_mode((w, h), flags)
    except Exception:
        screen = pygame.display.set_mode((w, h),
                                         pygame.FULLSCREEN if want_fs else 0)
    pygame.display.set_caption(caption)
    pygame.mouse.set_visible(False)
    return screen


class GKInput:
    """Pointer + motion for Fit games. pos is absolute screen coordinates."""

    def __init__(self, screen_w, screen_h, span_scale=1.0):
        self.W = screen_w
        self.H = screen_h
        self.conf = load_conf()
        self.span_scale = span_scale
        self.imu = None
        self.btns = []
        self.side = None
        self.res = {"g": 14247.0, "a": 4096.0}
        self.g = {"x": 0.0, "y": 0.0, "z": 0.0}
        self.acc = [0.0, 0.0, 1.0]
        self.bias = [0.0, 0.0, 0.0]
        self.yaw = 0.0
        self.pitch = 0.0
        self.pos = [screen_w / 2.0, screen_h / 2.0]
        self.pressed = set()        # buttons that went down this frame
        self.held = set()
        self.swing = 0.0            # peak deg/s over the last moment
        self._swing_decay = 0.0
        self.calibrating = 45
        self._cal = [0.0, 0.0, 0.0]
        self._cal_lo = [9e9, 9e9, 9e9]
        self._cal_hi = [-9e9, -9e9, -9e9]
        self._cal_tries = 3
        self._cal_frames = 45
        self.cal_status = ""
        self._cal_n = 0
        self._last_ts = None
        self._still = 0.0
        self._next_scan = 0.0
        self._sm = None
        self.pad = None
        self._pad_prev = {}
        self._mouse_ts = 0.0
        self.quit = False
        self._scan()
        self._init_pad()

    # ------------------------------------------------------------- setup
    @property
    def has_gyro(self):
        return self.imu is not None

    def _init_pad(self):
        try:
            pygame.joystick.init()
            if pygame.joystick.get_count():
                self.pad = pygame.joystick.Joystick(0)
                self.pad.init()
        except Exception:
            self.pad = None

    def _open(self, path):
        d = evdev.InputDevice(path)
        try:
            os.set_blocking(d.fd, False)
        except Exception:
            pass
        return d

    def _scan(self):
        self._next_scan = time.time() + 3.0
        if not HAVE_EVDEV:
            return
        try:
            paths = evdev.list_devices()
        except Exception:
            return
        have = {d.path for d in self.btns}
        imus = []
        for p in paths:
            try:
                d = self._open(p)
            except Exception:
                continue
            name = d.name or ""
            if "Joy-Con" not in name and "Nintendo" not in name:
                try:
                    d.close()
                except Exception:
                    pass
                continue
            if "IMU" in name:
                imus.append(d)
            elif p not in have:
                self.btns.append(d)
            else:
                try:
                    d.close()
                except Exception:
                    pass
        if self.imu is None and imus:
            self.imu = imus.pop(0)
            n = self.imu.name or ""
            self.side = "R" if "Right" in n else ("L" if "Left" in n else "?")

            def _res(code, dflt):
                try:
                    r = float(self.imu.absinfo(code).resolution)
                    return r if r > 0 else dflt
                except Exception:
                    return dflt
            self.res = {"g": _res(ecodes.ABS_RX, 14247.0),
                        "a": _res(ecodes.ABS_X, 4096.0)}
        for d in imus:
            try:
                d.close()
            except Exception:
                pass

    def recenter(self):
        """Instant re-aim. Doesn't touch the bias, so it can't go wrong."""
        self.yaw = 0.0
        self.pitch = 0.0
        self._sm = None

    def calibrate(self, frames=60, tries=3):
        """Learn the zero-rate bias, rejecting any window where it moved."""
        self._cal = [0.0, 0.0, 0.0]
        self._cal_lo = [9e9, 9e9, 9e9]
        self._cal_hi = [-9e9, -9e9, -9e9]
        self._cal_n = 0
        self._cal_tries = tries
        self._cal_frames = frames
        self.calibrating = frames
        self.cal_status = "working"
        self.recenter()
        self._last_ts = None

    # ------------------------------------------------------------- update
    def update(self, dt):
        self.pressed = set()
        for e in pygame.event.get():
            if e.type == pygame.QUIT:
                self.quit = True
            elif e.type == pygame.KEYDOWN:
                if e.key == pygame.K_ESCAPE:
                    self.quit = True
                elif e.key in (pygame.K_RETURN, pygame.K_SPACE):
                    self.pressed.add("a")
                    self.held.add("a")
                elif e.key == pygame.K_BACKSPACE:
                    self.pressed.add("b")
                elif e.key == pygame.K_r:
                    self.recenter()
            elif e.type == pygame.KEYUP:
                if e.key in (pygame.K_RETURN, pygame.K_SPACE):
                    self.held.discard("a")
            elif e.type == pygame.MOUSEMOTION:
                self.pos = list(e.pos)
                self._mouse_ts = time.time()
            elif e.type == pygame.MOUSEBUTTONDOWN and e.button == 1:
                self.pos = list(e.pos)
                self._mouse_ts = time.time()
                self.pressed.add("a")
                self.held.add("a")
            elif e.type == pygame.MOUSEBUTTONUP and e.button == 1:
                self.held.discard("a")

        self._read_imu(dt)
        self._read_buttons()
        self._read_pad()

        if "y" in self.pressed:
            self.recenter()
        self.pos[0] = max(0, min(self.W - 1, self.pos[0]))
        self.pos[1] = max(0, min(self.H - 1, self.pos[1]))
        self._swing_decay = max(0.0, self._swing_decay - dt * 2.5)
        self.swing *= (1.0 - min(1.0, dt * 3.0))

    def _read_imu(self, dt):
        if self.imu is None:
            if time.time() >= self._next_scan:
                self._scan()
            return
        samples = []
        try:
            while True:
                e = self.imu.read_one()
                if e is None:
                    break
                if e.type == ecodes.EV_ABS:
                    if e.code == ecodes.ABS_RX:
                        self.g["x"] = e.value / self.res["g"]
                    elif e.code == ecodes.ABS_RY:
                        self.g["y"] = e.value / self.res["g"]
                    elif e.code == ecodes.ABS_RZ:
                        self.g["z"] = e.value / self.res["g"]
                    elif e.code == ecodes.ABS_X:
                        self.acc[0] = e.value / self.res["a"]
                    elif e.code == ecodes.ABS_Y:
                        self.acc[1] = e.value / self.res["a"]
                    elif e.code == ecodes.ABS_Z:
                        self.acc[2] = e.value / self.res["a"]
                elif e.type == ecodes.EV_SYN and e.code == ecodes.SYN_REPORT:
                    ts = e.timestamp()
                    sdt = 0.005 if self._last_ts is None else ts - self._last_ts
                    if not (0.0 < sdt < 0.1):
                        sdt = 0.005
                    self._last_ts = ts
                    samples.append((self.g["x"], self.g["y"], self.g["z"], sdt))
        except (OSError, IOError):
            try:
                self.imu.close()
            except Exception:
                pass
            self.imu = None
            self._last_ts = None
            return
        except Exception:
            pass

        if not samples:
            return

        if self.calibrating > 0:
            for gx, gy, gz, _ in samples:
                for i, v in enumerate((gx, gy, gz)):
                    self._cal[i] += v
                    self._cal_lo[i] = min(self._cal_lo[i], v)
                    self._cal_hi[i] = max(self._cal_hi[i], v)
                self._cal_n += 1
            self.calibrating -= 1
            if self.calibrating <= 0:
                if self._cal_n >= 8:
                    spread = max(self._cal_hi[i] - self._cal_lo[i]
                                 for i in range(3))
                    if spread <= 2.5:
                        self.bias = [a / self._cal_n for a in self._cal]
                        self.cal_status = "ok"
                    elif self._cal_tries > 1:
                        self.calibrate(self._cal_frames, self._cal_tries - 1)
                    else:
                        self.cal_status = "gaveup"
                self.recenter()
            return

        conf = self.conf
        ya, pa = GYRO_MAPS[conf.get("gyro_map", 0) % len(GYRO_MAPS)]
        peak = 0.0
        for gx, gy, gz, sdt in samples:
            gx -= self.bias[0]
            gy -= self.bias[1]
            gz -= self.bias[2]
            axes = {"x": gx, "y": gy, "z": gz}
            yr, pr = axes[ya], axes[pa]
            if conf.get("swap_axes"):
                yr, pr = pr, yr
            mag = math.sqrt(gx * gx + gy * gy + gz * gz)
            peak = max(peak, mag)
            if abs(yr) > 0.35:
                self.yaw += yr * sdt
            if abs(pr) > 0.35:
                self.pitch += pr * sdt

        self.swing = max(self.swing, peak)

        if peak < 2.2:
            self._still += dt
            if self._still > 0.25:
                raw = [sum(s[i] for s in samples) / len(samples)
                       for i in range(3)]
                k = min(1.0, dt * (2.2 if self._still > 1.0 else 0.9))
                for i in range(3):
                    self.bias[i] += (raw[i] - self.bias[i]) * k
                self.cal_status = "ok"
        else:
            self._still = 0.0

        span = max(12.0, 46.0 / max(0.3, conf.get("sens", 1.0)) * self.span_scale)
        half = span * 0.5
        self.yaw = max(-half, min(half, self.yaw))
        self.pitch = max(-half, min(half, self.pitch))
        ppd = self.W / span
        tx = self.W * 0.5 - self.yaw * ppd
        ty = self.H * 0.5 - self.pitch * ppd
        if conf.get("invert_x"):
            tx = self.W - tx
        if conf.get("invert_y"):
            ty = self.H - ty
        if self._sm is None:
            self._sm = [tx, ty]
        k = min(1.0, dt * 26.0)
        self._sm[0] += (tx - self._sm[0]) * k
        self._sm[1] += (ty - self._sm[1]) * k
        if time.time() - self._mouse_ts > 0.15:
            self.pos = [self._sm[0], self._sm[1]]

    def _read_buttons(self):
        dead = []
        for d in self.btns:
            try:
                while True:
                    e = d.read_one()
                    if e is None:
                        break
                    if e.type == ecodes.EV_KEY:
                        name = BTN_MAP.get(e.code)
                        if not name:
                            continue
                        if e.value == 1:
                            self.pressed.add(name)
                            self.held.add(name)
                        elif e.value == 0:
                            self.held.discard(name)
            except (OSError, IOError):
                dead.append(d)
            except Exception:
                pass
        for d in dead:
            try:
                d.close()
            except Exception:
                pass
            self.btns.remove(d)

    def _read_pad(self):
        if not self.pad:
            return
        if self.conf.get("pad_layout", "nintendo") == "nintendo":
            a_i, b_i = 1, 0
        else:
            a_i, b_i = 0, 1
        for name, idx in (("a", a_i), ("b", b_i), ("y", 3), ("x", 2)):
            try:
                v = self.pad.get_button(idx)
            except Exception:
                v = 0
            if v and not self._pad_prev.get(name):
                self.pressed.add(name)
            if v:
                self.held.add(name)
            else:
                self.held.discard(name)
            self._pad_prev[name] = v
        if not self.has_gyro:
            try:
                ax, ay = self.pad.get_axis(0), self.pad.get_axis(1)
            except Exception:
                return
            dead = 0.22
            if abs(ax) > dead:
                self.pos[0] += ax * self.W * 0.012
            if abs(ay) > dead:
                self.pos[1] += ay * self.W * 0.012

    # ------------------------------------------------------------- extras
    def tilt(self):
        """(-1..1, -1..1) from the accelerometer, for steering games."""
        ax, ay, az = self.acc
        n = math.hypot(ax, ay, az)
        if n < 0.25:
            return (self.pos[0] / self.W - 0.5) * 2, (self.pos[1] / self.H - 0.5) * 2
        return max(-1, min(1, ax / n * 2.2)), max(-1, min(1, ay / n * 2.2))

    def status(self):
        if self.has_gyro:
            return "Joy-Con (%s)  -  point at the screen" % self.side
        return "No Joy-Con  -  mouse / pad"