#!/usr/bin/env python3
import os
import sys
import time
import json
import re
import traceback
from typing import Any, Dict, Tuple, List, Optional

try:
    import pygame  # type: ignore
except Exception as e:
    print("Pygame is required. Install with: pip install pygame")
    raise


PROMPTS = [
    # Face buttons
    "SOUTH_BUTTON - CONFIRM",  # A on Xbox
    "EAST_BUTTON - CANCEL",   # B on Xbox
    "WEST_BUTTON - CLEAR HISTORY / SELECT GAMES",   # X on Xbox
    "NORTH_BUTTON - HISTORY",  # Y on Xbox
    # Meta
    "START - PAUSE",
    "SELECT - FILTER",
    # D-Pad
    "DPAD_UP - MOVE UP",
    "DPAD_DOWN - MOVE DOWN",
    "DPAD_LEFT - MOVE LEFT",
    "DPAD_RIGHT - MOVE  RIGHT",
    # Bumpers
    "LEFT_BUMPER - LB/L1 - Delete last char",
    "RIGHT_BUMPER - RB/R1 - Add space",
    # Triggers
    "LEFT_TRIGGER - LT/L2 - Page +",
    "RIGHT_TRIGGER - RT/R2 - Page -",
    # Left stick directions
    "JOYSTICK_LEFT_UP - MOVE UP",
    "JOYSTICK_LEFT_DOWN - MOVE DOWN",
    "JOYSTICK_LEFT_LEFT - MOVE LEFT",
    "JOYSTICK_LEFT_RIGHT - MOVE RIGHT",
]

INPUT_TIMEOUT_SECONDS = 10  # Temps max par entrée avant "ignored"

# --- Minimal on-screen console (Pygame window) ---
SURFACE = None  # type: ignore
FONT = None  # type: ignore
LOG_LINES: List[str] = []
MAX_LOG = 300


def init_screen(width: int = 900, height: int = 600) -> None:
    global SURFACE, FONT
    try:
        pygame.display.init()
        SURFACE = pygame.display.set_mode((width, height))
        pygame.display.set_caption("Controller Tester")
        pygame.font.init()
        FONT = pygame.font.SysFont("Consolas", 20) or pygame.font.Font(None, 20)
    except Exception:
        # If display init fails, stay headless but continue
        SURFACE = None
        FONT = None


def log(msg: str) -> None:
    # Print to real console and on-screen log
    try:
        print(msg)
    except Exception:
        pass
    LOG_LINES.append(str(msg))
    if len(LOG_LINES) > MAX_LOG:
        del LOG_LINES[: len(LOG_LINES) - MAX_LOG]
    draw_log()


def draw_log() -> None:
    if SURFACE is None or FONT is None:
        return
    try:
        SURFACE.fill((12, 12, 12))
        margin = 12
        line_h = FONT.get_height() + 4
        # Show the last N lines that fit on screen
        max_lines = (SURFACE.get_height() - margin * 2) // line_h
        to_draw = LOG_LINES[-max_lines:]
        y = margin
        for line in to_draw:
            surf = FONT.render(line, True, (220, 220, 220))
            SURFACE.blit(surf, (margin, y))
            y += line_h
        pygame.display.flip()
    except Exception:
        pass


def init_joystick() -> pygame.joystick.Joystick:
    pygame.init()
    pygame.joystick.init()
    if pygame.joystick.get_count() == 0:
        log("No joystick detected. Connect a controller and try again.")
        sys.exit(1)
    js = pygame.joystick.Joystick(0)
    js.init()
    name = js.get_name()
    log(f"Using joystick 0: {name}")
    log("")
    log(f"Note: each input will auto-ignore after {INPUT_TIMEOUT_SECONDS}s if not present (e.g. missing L2/R2)")
    return js


def wait_for_stable(js: pygame.joystick.Joystick, settle_ms: int = 250, deadband: float = 0.05, timeout_ms: int = 2000) -> bool:
    """Wait until axes stop moving (change < deadband) continuously for settle_ms.

    Unlike a traditional neutral check, this doesn't assume axes center at 0.
    Hats are required to be (0,0) to avoid capturing D-Pad releases.
    Returns True if stability achieved, False on timeout.
    """
    start = pygame.time.get_ticks()
    last = [js.get_axis(i) for i in range(js.get_numaxes())]
    stable_since = None
    while True:
        # Handle window close only (avoid quitting on keyboard here)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                log("Window closed. Exiting.")
                sys.exit(0)
        moved = False
        for i in range(js.get_numaxes()):
            cur = js.get_axis(i)
            if abs(cur - last[i]) > deadband:
                moved = True
            last[i] = cur
        hats_ok = all(js.get_hat(i) == (0, 0) for i in range(js.get_numhats()))
        if not moved and hats_ok:
            if stable_since is None:
                stable_since = pygame.time.get_ticks()
            elif pygame.time.get_ticks() - stable_since >= settle_ms:
                return True
        else:
            stable_since = None
        if pygame.time.get_ticks() - start > timeout_ms:
            return False
        draw_log()
        pygame.time.wait(10)


def wait_for_event(js: pygame.joystick.Joystick, logical_name: str, axis_threshold: float = 0.6, timeout_sec: int = INPUT_TIMEOUT_SECONDS) -> Tuple[str, Any]:
    """Wait for a joystick event for the given logical control.

    Returns a tuple of (kind, data):
      - ("button", button_index)
      - ("hat", (x, y)) where x,y in {-1,0,1}
      - ("axis", {"axis": index, "direction": -1|1})
    """
    # Ensure prior motion has settled to avoid capturing a release
    wait_for_stable(js)
    log("")
    deadline = time.time() + max(1, int(timeout_sec))
    log(f"Press {logical_name} (Wait {timeout_sec}s to skip/ignore) if not present")
    # Flush old events
    pygame.event.clear()
    while True:
        # Update window title with countdown if we have a surface
        try:
            remaining = int(max(0, deadline - time.time()))
            if SURFACE is not None:
                pygame.display.set_caption(f"Controller Tester — {logical_name} — {remaining}s left")
        except Exception:
            pass
        for event in pygame.event.get():
            # Keyboard helpers
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    log(f"Skipped {logical_name}")
                    return ("skipped", None)
                # No keyboard quit here to avoid accidental exits when using controllers
            if event.type == pygame.QUIT:
                log("Window closed. Exiting.")
                sys.exit(0)

            # Buttons
            if event.type == pygame.JOYBUTTONDOWN:
                log(f"Captured {logical_name}: BUTTON {event.button}")
                return ("button", event.button)

            # D-Pad (HAT)
            if event.type == pygame.JOYHATMOTION:
                val = event.value  # (x, y)
                if val != (0, 0):
                    log(f"Captured {logical_name}: HAT {val}")
                    return ("hat", val)

            # Axes (sticks, triggers)
            if event.type == pygame.JOYAXISMOTION:
                axis = event.axis
                value = float(event.value)
                if abs(value) >= axis_threshold:
                    direction = 1 if value > 0 else -1
                    log(f"Captured {logical_name}: AXIS {axis} dir {direction} (raw {value:.2f})")
                    return ("axis", {"axis": axis, "direction": direction, "raw": value})

        draw_log()
        # Timeout?
        if time.time() >= deadline:
            log(f"Ignored {logical_name} (timeout {timeout_sec}s)")
            return ("ignored", None)
        time.sleep(0.005)


def write_log(path: str, mapping: Dict[str, Tuple[str, Any]], device_name: str) -> None:
    lines = []
    lines.append("# Controller mapping log\n")
    lines.append(f"# Device: {device_name}\n\n")
    for name, (kind, data) in mapping.items():
        if kind == "button":
            lines.append(f"{name} = BUTTON {data}\n")
        elif kind == "hat":
            lines.append(f"{name} = HAT {data}\n")
        elif kind == "axis":
            ax = data.get("axis")
            direction = data.get("direction")
            lines.append(f"{name} = AXIS {ax} dir {direction}\n")
        elif kind == "skipped":
            lines.append(f"{name} = SKIPPED\n")
        elif kind == "ignored":
            lines.append(f"{name} = IGNORED\n")
        else:
            lines.append(f"{name} = UNKNOWN {data}\n")

    with open(path, "w", encoding="utf-8") as f:
        f.writelines(lines)
    log("")
    log(f"Saved mapping to: {path}")


# --- JSON preset generation ---
def sanitize_device_name(name: str) -> str:
    s = name.strip().lower()
    # Replace non-alphanumeric with underscore
    s = re.sub(r"[^a-z0-9]+", "_", s)
    s = re.sub(r"_+", "_", s).strip("_")
    return s or "controller"


def to_json_binding(kind: str, data: Any, display: Optional[str] = None) -> Optional[Dict[str, Any]]:
    if kind == "button" and isinstance(data, int):
        return {"type": "button", "button": data, **({"display": display} if display else {})}
    if kind == "hat" and isinstance(data, (tuple, list)) and len(data) == 2:
        val = list(data)
        return {"type": "hat", "value": val, **({"display": display} if display else {})}
    if kind == "axis" and isinstance(data, dict):
        axis = data.get("axis")
        direction = data.get("direction")
        if isinstance(axis, int) and direction in (-1, 1):
            return {"type": "axis", "axis": axis, "direction": int(direction), **({"display": display} if display else {})}
    return None


def build_controls_json(mapping: Dict[str, Tuple[str, Any]]) -> Dict[str, Any]:
    # Map logical prompts to action keys and preferred display labels
    prompt_map = {
        "SOUTH_BUTTON - CONFIRM": ("confirm", "A"),
        "EAST_BUTTON - CANCEL": ("cancel", "B"),
        "WEST_BUTTON - CLEAR HISTORY / SELECT GAMES": ("clear_history", "X"),
        "NORTH_BUTTON - HISTORY": ("history", "Y"),
        "START - PAUSE": ("start", "Start"),
        "SELECT - FILTER": ("filter", "Select"),
        "DPAD_UP - MOVE UP": ("up", "↑"),
        "DPAD_DOWN - MOVE DOWN": ("down", "↓"),
        "DPAD_LEFT - MOVE LEFT": ("left", "←"),
        "DPAD_RIGHT - MOVE  RIGHT": ("right", "→"),
        "LEFT_BUMPER - LB/L1 - Delete last char": ("delete", "LB"),
        "RIGHT_BUMPER - RB/R1 - Add space": ("space", "RB"),
        # Triggers per prompts: LEFT=page_up, RIGHT=page_down
        "LEFT_TRIGGER - LT/L2 - Page +": ("page_up", "LT"),
        "RIGHT_TRIGGER - RT/R2 - Page -": ("page_down", "RT"),
        # Left stick directions (fallbacks for arrows)
        "JOYSTICK_LEFT_UP - MOVE UP": ("up", "J↑"),
        "JOYSTICK_LEFT_DOWN - MOVE DOWN": ("down", "J↓"),
        "JOYSTICK_LEFT_LEFT - MOVE LEFT": ("left", "J←"),
        "JOYSTICK_LEFT_RIGHT - MOVE RIGHT": ("right", "J→"),
    }

    result: Dict[str, Any] = {}

    # First pass: take direct DPAD/face/meta/bumper/trigger bindings
    for prompt, (action, disp) in prompt_map.items():
        if prompt not in mapping:
            continue
        kind, data = mapping[prompt]
        if kind in ("ignored", "skipped"):
            continue
        # Prefer DPAD over JOYSTICK for directions: handle fallback later
        if action in ("up", "down", "left", "right"):
            if prompt.startswith("DPAD_"):
                b = to_json_binding(kind, data, disp)
                if b:
                    result[action] = b
            # Joystick handled as fallback if DPAD missing
        else:
            b = to_json_binding(kind, data, disp)
            if b:
                result[action] = b

    # Second pass: fallback to joystick directions if arrows missing
    fallbacks = [
        ("JOYSTICK_LEFT_UP - MOVE UP", "up", "J↑"),
        ("JOYSTICK_LEFT_DOWN - MOVE DOWN", "down", "J↓"),
        ("JOYSTICK_LEFT_LEFT - MOVE LEFT", "left", "J←"),
        ("JOYSTICK_LEFT_RIGHT - MOVE RIGHT", "right", "J→"),
    ]
    for prompt, action, disp in fallbacks:
        if action in result:
            continue
        if prompt in mapping:
            kind, data = mapping[prompt]
            if kind in ("ignored", "skipped"):
                continue
            b = to_json_binding(kind, data, disp)
            if b:
                result[action] = b

    return result


def write_controls_json(device_name: str, controls: Dict[str, Any]) -> str:
    """Write the generated controls preset JSON in the same folder as this script.

    Also embeds a JSON-safe comment with the device name under the _comment key.
    """
    # Same folder as the launched script
    base_dir = os.path.dirname(os.path.abspath(__file__))
    fname = f"{sanitize_device_name(device_name)}_controller.json"
    out_path = os.path.join(base_dir, fname)
    # Include the detected device name for auto-preset matching
    payload = {"device": device_name}
    payload.update(controls)
    try:
        with open(out_path, "w", encoding="utf-8") as f:
            json.dump(payload, f, ensure_ascii=False, indent=4)
        return out_path
    except Exception:
        return out_path


def main() -> None:
    init_screen()
    js = init_joystick()
    # Print device basics
    try:
        log(f"Buttons: {js.get_numbuttons()} | Axes: {js.get_numaxes()} | Hats: {js.get_numhats()}")
    except Exception:
        pass

    mapping: Dict[str, Tuple[str, Any]] = {}
    for logical in PROMPTS:
        kind, data = wait_for_event(js, logical)
        mapping[logical] = (kind, data)
        # Short, consistent debounce for all inputs
        pygame.event.clear()
        pygame.time.wait(150)

    log_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "controller_mapping.log")
    write_log(log_path, mapping, js.get_name())
    # Build and write ready-to-use JSON controls preset
    controls = build_controls_json(mapping)
    if controls:
        out_json = write_controls_json(js.get_name(), controls)
        log(f"Saved JSON preset to: {out_json}")
    else:
        log("No usable inputs captured to build a JSON preset.")
    log("Done. Press Q or close the window to exit.")


if __name__ == "__main__":
    try:
        main()
    except SystemExit:
        # Allow intentional exits
        pass
    except Exception:
        # Show traceback on screen and wait for window close
        tb = traceback.format_exc()
        try:
            log("")
            log("An error occurred:")
            for line in tb.splitlines():
                log(line)
            # Idle until window is closed
            while True:
                for event in pygame.event.get():
                    if event.type == pygame.QUIT:
                        raise SystemExit(1)
                draw_log()
                pygame.time.wait(50)
        except Exception:
            pass
    finally:
        try:
            pygame.joystick.quit()
            pygame.quit()
        except Exception:
            pass