# -*- coding: utf-8 -*-
"""Keyboard teleoperation goal using an input listener separate from MuJoCo."""

import atexit
import queue
import threading

from robogauge.tasks.gauge.goal_data import GoalData, VelocityGoal
from robogauge.tasks.gauge.goals.base_goal import BaseGoal
from robogauge.tasks.simulator.sim_data import SimData
from robogauge.utils.helpers import class_to_dict
from robogauge.utils.logger import logger


class _KeyboardState:
    """Track held keys without replacing MuJoCo viewer shortcuts."""

    def __init__(self):
        self.events = queue.Queue()
        self.running = False
        self.held = set()
        self.lock = threading.Lock()
        self.thread = None
        self.listener = None

    @staticmethod
    def _name(key):
        try:
            if hasattr(key, 'char') and key.char:
                return key.char.lower()
        except Exception:
            pass
        return str(key).lower()

    def _worker(self):
        while self.running:
            try:
                event_type, key = self.events.get(timeout=0.05)
            except queue.Empty:
                continue
            name = self._name(key)
            with self.lock:
                if event_type == 'press':
                    self.held.add(name)
                else:
                    self.held.discard(name)

    def start(self):
        try:
            from pynput import keyboard
        except ImportError as exc:
            raise ImportError(
                "Keyboard control requires 'pynput'. Install it with "
                "'python -m pip install pynput'."
            ) from exc

        self.running = True
        self.listener = keyboard.Listener(
            on_press=lambda key: self.events.put(('press', key)),
            on_release=lambda key: self.events.put(('release', key)),
        )
        self.listener.start()
        self.thread = threading.Thread(target=self._worker, daemon=True)
        self.thread.start()

    def snapshot(self):
        with self.lock:
            return set(self.held)

    def clear(self):
        with self.lock:
            self.held.clear()

    def stop(self):
        self.running = False
        if self.listener is not None:
            self.listener.stop()
            self.listener = None


class KeyboardGoal(BaseGoal):
    """Produce velocity commands while the configured teleoperation keys are held."""

    name = 'keyboard'

    def __init__(self, max_velocity, goal_type='velocity', **kwargs):
        super().__init__()
        if goal_type != 'velocity':
            raise NotImplementedError("Only 'velocity' goal type is implemented for KeyboardGoal.")

        kwargs.pop('enabled', None)
        if kwargs:
            logger.warning(f"Unused kwargs in KeyboardGoal: {kwargs}")

        self.max_velocity = class_to_dict(max_velocity)
        self.total = 1
        self.current_command = (0.0, 0.0, 0.0)
        self.keyboard = _KeyboardState()
        self.keyboard.start()
        atexit.register(self.keyboard.stop)
        logger.info(
            "Keyboard control: Up/Down=forward/back, Left/Right=yaw, "
            "','=left strafe, '.'=right strafe, K=stop. "
            "MuJoCo viewer shortcuts remain independent."
        )

    def is_reset(self, sim_data: SimData) -> bool:
        return False

    def pre_get_goal(self, sim_data: SimData) -> bool:
        return False

    def reset_goal(self):
        self.keyboard.clear()
        self.current_command = (0.0, 0.0, 0.0)

    @staticmethod
    def _axis_command(held, positive_key, negative_key, limits):
        positive = positive_key in held
        negative = negative_key in held
        if positive == negative:
            return 0.0
        return float(limits[1] if positive else limits[0])

    def command_from_keys(self, held):
        if 'k' in held:
            return 0.0, 0.0, 0.0
        return (
            self._axis_command(held, 'key.up', 'key.down', self.max_velocity['lin_vel_x']),
            self._axis_command(held, ',', '.', self.max_velocity['lin_vel_y']),
            self._axis_command(held, 'key.left', 'key.right', self.max_velocity['ang_vel_yaw']),
        )

    def get_goal(self, sim_data: SimData) -> GoalData:
        command = self.command_from_keys(self.keyboard.snapshot())
        if command != self.current_command:
            self.current_command = command
            logger.info(
                "Keyboard command: vx=%.2f, vy=%.2f, yaw=%.2f",
                *command,
            )
        return GoalData(
            goal_type='velocity',
            velocity_goal=VelocityGoal(
                lin_vel_x=command[0],
                lin_vel_y=command[1],
                ang_vel_yaw=command[2],
            ),
        )