{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE NamedFieldPuns #-}

module Twirl.Inputs (
  InputState
  , ControllerAxis (..)
  , ControllerButton (..)
  , MouseButton (..)
  , updateInputs
  , closeControllers
  , emptyInputState
  , numControllers
  , isControllerButtonDown
  , controllerAxis
  , isKeyDown
  , isMouseDown
  , mousePosition
) where

import Control.Monad.IO.Class (MonadIO)
import Data.Foldable (foldlM)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe)
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Word (Word8)
import Foreign.C.Types (CInt)
import GHC.Int (Int16, Int32)
import SDL (
  ControllerAxisEventData (..)
  , ControllerButtonEventData (..)
  , ControllerDeviceEventData (..)
  , Event (..)
  , EventPayload (ControllerAxisEvent, ControllerButtonEvent, ControllerDeviceEvent, KeyboardEvent, QuitEvent)
  , InputMotion (Pressed, Released)
  , KeyboardEventData (..)
  , Keycode
  , Keysym (..)
  , MouseButton (..)
  , Point (..)
  , V2 (..)
  , getAbsoluteMouseLocation
  , getMouseButtons
  , keyboardEventKeyMotion
  , pollEvents
 )
import SDL.Input.GameController (
  ControllerButton (..)
  , ControllerButtonState (..)
  , ControllerDeviceConnection (..)
 )
import SDL.Raw.Event (gameControllerClose, gameControllerFromInstanceID, gameControllerGetJoystick, gameControllerOpen, joystickInstanceID)
import SDL.Raw.Types (GameController)
import Twirl.Graphics

data ControllerAxis = ControllerAxisLeftX | ControllerAxisLeftY | ControllerAxisRightX | ControllerAxisRightY | ControllerAxisTriggerLeft | ControllerAxisTriggerRight

data ControllerAxes = ControllerAxes
  { controllerAxisLeftX :: Int16
  , controllerAxisLeftY :: Int16
  , controllerAxisRightX :: Int16
  , controllerAxisRightY :: Int16
  , controllerAxisTriggerLeft :: Int16
  , controllerAxisTriggerRight :: Int16
  }
  deriving (Show)

data ControllerState = ControllerState
  { controllerStateButtons :: Set ControllerButton
  , controllerStateAxes :: ControllerAxes
  }
  deriving (Show)

emptyControllerState :: ControllerState
emptyControllerState =
  ControllerState
    { controllerStateButtons = Set.empty
    , controllerStateAxes =
        ControllerAxes
          { controllerAxisLeftX = 0
          , controllerAxisLeftY = 0
          , controllerAxisRightX = 0
          , controllerAxisRightY = 0
          , controllerAxisTriggerLeft = 0
          , controllerAxisTriggerRight = 0
          }
    }

data InputState = InputState
  { controllers :: Map Int ControllerState
  , keyDownSet :: Set Keycode
  , instanceIdToControllerNo :: Map Int32 Int
  , mouseButtons :: (MouseButton -> Bool)
  , mousePos :: (Int, Int)
  }
  deriving (Show)

instance Show (MouseButton -> Bool) where
  show _ = "<mouseButtonFunction>"

updateInputs :: MonadIO m => InputState -> m (InputState, Bool)
updateInputs inputState = do
  events <- pollEvents
  mouseInputFn <- getMouseButtons
  mouseLocation <- getAbsoluteMouseLocation
  let quitApp = any ((== QuitEvent) . eventPayload) events
  updatedInputs <- handleEvents events inputState{mouseButtons = mouseInputFn, mousePos = toMousePos mouseLocation}
  pure (updatedInputs, quitApp)

toMousePos :: Point V2 CInt -> (Int, Int)
toMousePos (P (V2 x y)) = (fromIntegral x, fromIntegral y)

emptyInputState :: InputState
emptyInputState = InputState{controllers = Map.empty, keyDownSet = Set.empty, instanceIdToControllerNo = Map.empty, mouseButtons = const False, mousePos = (0, 0)}

getControllerAxis :: ControllerAxis -> ControllerState -> Int16
getControllerAxis ControllerAxisLeftX controllerState = controllerAxisLeftX $ controllerStateAxes controllerState
getControllerAxis ControllerAxisLeftY controllerState = controllerAxisLeftY $ controllerStateAxes controllerState
getControllerAxis ControllerAxisRightX controllerState = controllerAxisRightX $ controllerStateAxes controllerState
getControllerAxis ControllerAxisRightY controllerState = controllerAxisRightY $ controllerStateAxes controllerState
getControllerAxis ControllerAxisTriggerLeft controllerState = controllerAxisTriggerLeft $ controllerStateAxes controllerState
getControllerAxis ControllerAxisTriggerRight controllerState = controllerAxisTriggerRight $ controllerStateAxes controllerState

-- | Get number of controllers
--
-- Returns the maximum number of controllers that have been connected simultaneously. This means that if
-- you connect four controllers and unplug one, this function will return 4.
numControllers :: InputState -> Int
numControllers InputState{controllers} = (length . Map.elems) controllers

-- | Query the state of controller buttons
--
-- Returns True if the queried button is down on the queried controller. Returns False if the controller is disconnected or doesn't exist
-- (e.g. you query controller number 4 but there are only 2 controllers).
isControllerButtonDown ::
  InputState ->
  -- | Controller number to be queried
  Int ->
  -- | The Controller button to be queried
  ControllerButton ->
  Bool
isControllerButtonDown InputState{controllers} controllerNumber button =
  case Map.lookup controllerNumber controllers of
    Nothing -> False
    Just controllerState -> Set.member button (controllerStateButtons controllerState)

-- | Query axes of controllers
--
-- For the stick axes, the returned value is between -1.0 and 1.0, where 0.0 is the neutral position.
-- For triggers the value is between 0.0 and 1.0 where 0.0 is the neutral position. Bear in mind that
-- usually the level of quality of controller sticks is so low that the sticks don't return exactly to
-- the 0.0 position, so you have to ignore some values which are too close to zero.
controllerAxis ::
  InputState ->
  -- | Controller number
  Int ->
  -- | The queried axis
  ControllerAxis ->
  Double
controllerAxis InputState{controllers} controllerNumber axis =
  case Map.lookup controllerNumber controllers of
    Nothing -> 0
    Just controllerState -> (fromIntegral $ getControllerAxis axis controllerState) / 32767

-- | Query keyboard state
isKeyDown :: InputState -> Keycode -> Bool
isKeyDown InputState{keyDownSet} keycode = Set.member keycode keyDownSet

-- | Query mouse button state
isMouseDown :: InputState -> MouseButton -> Bool
isMouseDown InputState{mouseButtons} = mouseButtons

-- | Get mouse position on screen
mousePosition :: InputState -> (Int, Int)
mousePosition InputState{mousePos} = mousePos

handleEvents :: MonadIO m => [Event] -> InputState -> m InputState
handleEvents events keySet =
  foldlM
    handleEvent
    keySet
    events

handleEvent :: MonadIO m => InputState -> Event -> m InputState
handleEvent inputState event =
  case eventPayload event of
    KeyboardEvent keyboardEvent ->
      case keyboardEventKeyMotion keyboardEvent of
        Pressed -> pure inputState{keyDownSet = Set.insert (keysymKeycode $ keyboardEventKeysym keyboardEvent) (keyDownSet inputState)}
        Released -> pure inputState{keyDownSet = Set.delete (keysymKeycode $ keyboardEventKeysym keyboardEvent) (keyDownSet inputState)}
    ControllerDeviceEvent (ControllerDeviceEventData{controllerDeviceEventConnection, controllerDeviceEventWhich}) -> case controllerDeviceEventConnection of
      -- Type cast is necessary because SDL returns either device index or instance id
      -- (https://wiki.libsdl.org/SDL_ControllerDeviceEvent) depending on whether the
      -- controller was added or removed, but Haskell bindings always report to return
      -- instance id (which must be wrong since added but unopened controllers don't have
      -- an instance id yet)
      ControllerDeviceAdded -> gameControllerOpen (fromIntegral controllerDeviceEventWhich) >>= addControllerMapping inputState
      ControllerDeviceRemoved -> do
        controllerPtr <- gameControllerFromInstanceID controllerDeviceEventWhich
        updatedInputs <- removeControllerMapping inputState controllerPtr
        gameControllerClose controllerPtr
        pure updatedInputs
      _ -> pure inputState
    ControllerButtonEvent e -> do
      let updatedState = updateInputStateController e inputState
      pure updatedState
    ControllerAxisEvent e -> do
      let updatedState = updateInputStateControllerAxis e inputState
      pure updatedState
    _ -> pure inputState

addControllerMapping :: MonadIO m => InputState -> GameController -> m InputState
addControllerMapping inputState@InputState{instanceIdToControllerNo} controller = do
  instanceId <- gameControllerGetJoystick controller >>= joystickInstanceID
  let controllerNo = nextFreeController inputState
  pure inputState{instanceIdToControllerNo = Map.insert instanceId controllerNo instanceIdToControllerNo}

nextFreeController :: InputState -> Int
nextFreeController InputState{instanceIdToControllerNo} =
  let mappedControllers = Set.fromList $ Map.elems instanceIdToControllerNo
   in findSmallestNaturalNotInSet mappedControllers
 where
  findSmallestNaturalNotInSet :: Set Int -> Int
  findSmallestNaturalNotInSet set =
    let go n = if not (Set.member n set) then n else go (n + 1)
     in go 0

removeControllerMapping :: MonadIO m => InputState -> GameController -> m InputState
removeControllerMapping inputState@InputState{instanceIdToControllerNo} controller = do
  instanceId <- gameControllerGetJoystick controller >>= joystickInstanceID
  pure inputState{instanceIdToControllerNo = Map.delete instanceId instanceIdToControllerNo}

closeControllers :: InputState -> TwirlMonad ()
closeControllers InputState{instanceIdToControllerNo} = do
  let controllerIds = Map.keys instanceIdToControllerNo
  controllerPtrs <- mapM gameControllerFromInstanceID controllerIds
  mapM_ gameControllerClose controllerPtrs

updateInputStateController :: ControllerButtonEventData -> InputState -> InputState
updateInputStateController ControllerButtonEventData{controllerButtonEventWhich, controllerButtonEventState, controllerButtonEventButton} inputState@InputState{controllers, instanceIdToControllerNo} =
  let controller = fromMaybe (nextFreeController inputState) $ Map.lookup controllerButtonEventWhich instanceIdToControllerNo
   in inputState{controllers = Map.alter (updateControllerState controllerButtonEventState controllerButtonEventButton) controller controllers}

updateControllerState :: ControllerButtonState -> ControllerButton -> Maybe ControllerState -> Maybe ControllerState
updateControllerState stateChange button maybeState =
  let state = fromMaybe emptyControllerState maybeState
      action = case stateChange of
        ControllerButtonPressed -> Set.insert
        ControllerButtonReleased -> Set.delete
        _ -> flip const
   in pure state{controllerStateButtons = action button (controllerStateButtons state)}

updateInputStateControllerAxis :: ControllerAxisEventData -> InputState -> InputState
updateInputStateControllerAxis ControllerAxisEventData{controllerAxisEventWhich, controllerAxisEventValue, controllerAxisEventAxis} inputState@InputState{controllers, instanceIdToControllerNo} =
  let controller = fromMaybe (nextFreeController inputState) $ Map.lookup controllerAxisEventWhich instanceIdToControllerNo
   in inputState{controllers = Map.alter (updateControllerStateAxis controllerAxisEventValue controllerAxisEventAxis) controller controllers}

updateControllerStateAxis :: Int16 -> Word8 -> Maybe ControllerState -> Maybe ControllerState
updateControllerStateAxis value axis maybeState =
  let state = fromMaybe emptyControllerState maybeState
   in pure state{controllerStateAxes = updateControllerAxis (controllerStateAxes state) value axis}

updateControllerAxis :: ControllerAxes -> Int16 -> Word8 -> ControllerAxes
updateControllerAxis axes value axis = case axis of
  0 -> axes{controllerAxisLeftX = value}
  1 -> axes{controllerAxisLeftY = value}
  2 -> axes{controllerAxisRightX = value}
  3 -> axes{controllerAxisRightY = value}
  4 -> axes{controllerAxisTriggerLeft = value}
  _ -> axes{controllerAxisTriggerRight = value}