from agxPythonModules.utils.callbacks import KeyboardCallback as Input
from openplx.Physics.Signals import RealInputSignal, IntInputSignal, BoolInputSignal


class VehicleManeuvering:
    """    
    CONTROL_MAP supports three row types:

    1) Pair-key with continuous hold (setpoint)
        Suitable where holding the key commands a constant value.
        For example a motion controlled hydraulic cylinder.

        (key_pos, key_neg, specs)
        specs: (input_attr, value_pos [, value_neg])
        - If value_neg omitted, it defaults to -value_pos.
        - While held: sends value_pos/value_neg
        - On release: sends 0
        
    2) Single-key row:
        Operates in three modes:
        - "toggle"     is for items that can be turned on/off (toggles on key press, ignores key-repeat)
        - "hold"       is for items that should be active as long as the key is pressed
        - "ramp"       is for incrementally increasing a value while the key is held down
                       (one increment per key-repeat event)
        - "ramp_reset" is for incrementally increasing a value while the key is held down
                       (one increment per key-repeat event, resets on release)

        (key, specs, mode)
        mode: "hold", "toggle", "ramp" or "ramp_reset"

        For mode == "hold" or "toggle":
            specs: (input_attr, on_value, off_value [, initial_value] )
            - initial_value defaults to False (0) if omitted

        For mode == "ramp":
            specs: (input_attr, step_per_event [, cap_value] [, initial_value])
            - step_per_event is added once per key-repeat event while held
            - the stepping is stopped when cap_value is reached (omit or pass None for no cap)
            - initial_value seeds the ramp state and is sent on initialization; to set it without
              a cap, pass None for the cap_value slot (if omitted or None, the ramp starts at 0.0)

        For mode == "ramp_reset":
            specs: (input_attr, step_per_event [, cap_value] [, reset_value] [, initial_value])
            - step_per_event is added once per key-repeat event while held
            - the stepping is stopped when cap_value is reached (omit or pass None for no cap)
            - reset_value defaults to 0.0 if omitted or None
            - initial_value seeds the ramp state and is sent on initialization; when specifying it,
              keep unused optional slots (cap/reset) filled with None (if omitted or None, the ramp
              starts at 0.0)

    3) Pair-key with stepped values (two keys, discrete step on key press):
        Suitable for items that are increased or decreased in discrete steps,
        e.g. changing gears or changing an angle setpoint.

        (key_inc, key_dec, specs, "step")
        specs: (input_attr, step, min_value, max_value [, initial_value])
        - Changes only on data.down (press), not on release.
        - Integer specs produce integer outputs.
        - Floating-point specs produce real outputs.
    """

    def __init__(self, input_queue, signal_interface, control_map, runtime_scalar_reader=None):
        super().__init__()
        self.input_queue = input_queue
        self.CONTROL_MAP = control_map
        self._runtime_scalar_reader = runtime_scalar_reader
        
        # State for toggle + step controls
        self._toggles = {}          # input_attr -> bool
        self._discrete = {}         # input_attr -> numeric step state
        self._ramps = {}  # input_attr -> float
        self._keys_held = set()     # tracks keys currently down

        inputs = signal_interface.getInputs()

        # Collect all unique input attribute names referenced by CONTROL_MAP
        input_attr_names = set()
        for row in self.CONTROL_MAP:
            kind = self._row_kind(row)
            if kind == "single":
                _, specs, _ = row
            elif kind == "pair_step":
                _, _, specs, _ = row
            else:  # "pair"
                _, _, specs = row

            for spec in specs:
                input_attr_names.add(spec[0])

        # Populate attributes dynamically
        for attr in input_attr_names:
            try:
                setattr(self, attr, inputs[attr])
            except KeyError as e:
                raise KeyError(f"Input '{attr}' not found in signal_interface inputs") from e

        # Initialize toggle states and push initial signal
        for row in self.CONTROL_MAP:
            kind = self._row_kind(row)
            if kind != "single":
                continue

            key, specs, mode = row
            if mode != "toggle":
                continue

            for spec in specs:
                attr, on_val, off_val, *rest = spec
                initial_state = bool(rest[0]) if rest else False

                # store state
                self._toggles[attr] = initial_state

                # send initial value so sim matches state from start
                self.send_signal_to(on_val if initial_state else off_val, getattr(self, attr))

        # Initialize discrete values
        for row in self.CONTROL_MAP:
            if self._row_kind(row) == "pair_step":
                _, _, specs, _ = row
                for spec in specs:
                    attr = spec[0]                    
                    vmin = spec[2]
                    vmax = spec[3]
                    initial = spec[4] if len(spec) >= 5 else vmin
                    is_integer = self._is_int_step_spec(spec)
                    value = max(vmin, min(vmax, initial))
                    self._discrete[attr] = int(value) if is_integer else float(value)

        # Initialize ramp values from optional initial_value slots
        for row in self.CONTROL_MAP:
            if self._row_kind(row) != "single":
                continue

            _, specs, mode = row
            if mode not in ("ramp", "ramp_reset"):
                continue

            for spec in specs:
                initial_index = 3 if mode == "ramp" else 4
                if len(spec) > initial_index:
                    initial_raw = spec[initial_index]
                    if initial_raw is not None:
                        initial_value = float(initial_raw)
                        attr = spec[0]
                        self._ramps[attr] = initial_value
                        self.send_signal_to(initial_value, getattr(self, attr))

    @staticmethod
    def _is_int_step_spec(spec) -> bool:
        return all(isinstance(value, int) and not isinstance(value, bool) for value in spec[1:])

    @staticmethod
    def _row_kind(row) -> str:
        # single: (key, specs, "hold"/"toggle"/"ramp"/"ramp_reset")
        if isinstance(row, tuple) and len(row) == 3 and isinstance(row[2], str):
            return "single"
        # pair_step: (key_inc, key_dec, specs, "step")
        if isinstance(row, tuple) and len(row) == 4 and row[3] == "step":
            return "pair_step"
        # pair: (key_pos, key_neg, specs)
        return "pair"

    def send_signal_to(self, value, input_signal):
        if isinstance(value, bool):
            self.input_queue.send(BoolInputSignal.create(bool(value), input_signal))
        elif isinstance(value, int) and not isinstance(value, bool):
            self.input_queue.send(IntInputSignal.create(int(value), input_signal))
        else:
            self.input_queue.send(RealInputSignal.create(float(value), input_signal))
            
    # Keyboard handling
    def handle_key(self, data) -> bool:
        for row in self.CONTROL_MAP:
            kind = self._row_kind(row)

            # Single-key row --------------------------------------------------
            if kind == "single":
                key, specs, mode = row
                if data.key != key:
                    continue

                if mode == "hold":
                    if data.down:
                        for attr, on_val, off_val in specs:
                            self.send_signal_to(on_val, getattr(self, attr))
                    else:
                        for attr, on_val, off_val in specs:
                            self.send_signal_to(off_val, getattr(self, attr))

                elif mode == "toggle":
                    if data.down:
                        # press-edge: ignore auto-repeat while held
                        if key in self._keys_held:
                            continue
                        self._keys_held.add(key)

                        for spec in specs:
                            attr, on_val, off_val, *rest = spec                        
                        # for attr, on_val, off_val in specs:
                            new_state = not self._toggles.get(attr, False)
                            self._toggles[attr] = new_state
                            self.send_signal_to(on_val if new_state else off_val, getattr(self, attr))
                            
                            # Optional: log actual state
                            print(f"TOGGLE {attr} -> {new_state}")                                                    
                    else:
                        # release: allow next press to toggle again
                        self._keys_held.discard(key)

                elif mode == "ramp" or mode == "ramp_reset":
                    if data.down:
                        # one increment per key-repeat event
                        for spec in specs:
                            attr = spec[0]
                            step = float(spec[1])
                            cur = float(self._ramps.get(attr, 0.0))
                            cur += step

                            # optionally cap the ramped value
                            cap_value = None
                            if len(spec) >= 3 and spec[2] is not None:
                                cap_value = float(spec[2])
                            if cap_value is not None:
                                cur = min(cur, cap_value) if step > 0 else max(cur, cap_value)

                            self._ramps[attr] = cur
                            self.send_signal_to(cur, getattr(self, attr))
                    else:
                        if mode == "ramp":
                            for spec in specs:
                                attr = spec[0]
                                if (
                                    attr not in ("cargo_bed_angle_input", "steering_angle_input")
                                    or self._runtime_scalar_reader is None
                                ):
                                    continue

                                current_value = self._runtime_scalar_reader(attr)
                                if current_value is None:
                                    continue

                                self._ramps[attr] = float(current_value)
                                self.send_signal_to(self._ramps[attr], getattr(self, attr))
                        elif mode == "ramp_reset":
                            # reset to 0 (or optional reset value)
                            for spec in specs:
                                attr = spec[0]
                                reset_val = spec[3] if len(spec) >= 4 else None
                                if reset_val is None:
                                    reset_val = 0.0
                                reset_val = float(reset_val)
                                self._ramps[attr] = reset_val
                                self.send_signal_to(reset_val, getattr(self, attr))

                else:
                    raise ValueError(f"Unknown single-key mode: {mode!r}")

                continue


            # Pair-step row ---------------------------------------------------
            if kind == "pair_step":
                key_inc, key_dec, specs, _ = row
                if data.key not in (key_inc, key_dec):
                    continue

                if data.down:
                    # If key is already held, it's an auto-repeat -> ignore
                    if data.key in self._keys_held:
                        continue
                    self._keys_held.add(data.key)

                    sign = +1 if data.key == key_inc else -1

                    for spec in specs:
                        attr = spec[0]
                        is_integer = self._is_int_step_spec(spec)
                        step = int(spec[1]) if is_integer else float(spec[1])
                        vmin = int(spec[2]) if is_integer else float(spec[2])
                        vmax = int(spec[3]) if is_integer else float(spec[3])

                        cur = self._discrete.get(attr, vmin)
                        nxt = max(vmin, min(vmax, cur + sign * step))

                        self._discrete[attr] = int(nxt) if is_integer else float(nxt)
                        self.send_signal_to(self._discrete[attr], getattr(self, attr))
                else:
                    # Key released -> allow the next press to trigger again
                    self._keys_held.discard(data.key)

                continue            

            # Pair row (continuous hold) --------------------------------------
            key_pos, key_neg, specs = row
            if data.key not in (key_pos, key_neg):
                continue

            # Track held keys
            if data.down:
                self._keys_held.add(data.key)
            else:
                self._keys_held.discard(data.key)

            # Determine commanded direction based on current held state
            pos_held = key_pos in self._keys_held
            neg_held = key_neg in self._keys_held

            # If neither held -> 0. If both held -> 0 (or pick a priority if you prefer)
            for spec in specs:
                attr = spec[0]
                value_pos = spec[1]
                value_neg = spec[2] if len(spec) >= 3 else -value_pos

                if pos_held and not neg_held:
                    cmd = value_pos
                elif neg_held and not pos_held:
                    cmd = value_neg
                else:
                    cmd = 0.0

                self.send_signal_to(cmd, getattr(self, attr))

        return False

    def setup_keyboard_listener(self):
        """
        - Pair row: <base>_pos / <base>_neg  (per input attr)
        - Single row: <base>_<mode>          (per input attr)
        - Pair-step row: <base>_inc / <base>_dec (per input attr)
        """
        for row in self.CONTROL_MAP:
            kind = self._row_kind(row)

            if kind == "single":
                key, specs, mode = row
                for spec in specs:
                    base = spec[0]
                    Input.bind(name=f"{base}_{mode}_{key}", key=key, callback=self.handle_key)

            elif kind == "pair_step":
                key_inc, key_dec, specs, _ = row
                for spec in specs:
                    base = spec[0]
                    Input.bind(name=f"{base}_inc_{key_inc}", key=key_inc, callback=self.handle_key)
                    Input.bind(name=f"{base}_dec_{key_dec}", key=key_dec, callback=self.handle_key)

            else:  # pair
                key_pos, key_neg, specs = row
                for spec in specs:
                    base = spec[0]
                    Input.bind(name=f"{base}_pos_{key_pos}", key=key_pos, callback=self.handle_key)
                    Input.bind(name=f"{base}_neg_{key_neg}", key=key_neg, callback=self.handle_key)
