{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedStrings #-}

module Hsm.Drive
  ( Speed(..)
  , Direction(..)
  , Duration
  , DriveAction(..)
  , DriveEnv(..)
  , driveEnvDefault
  , drive
  ) where

import Control.Concurrent (threadDelay)
import Control.Exception (AsyncException, handle)
import Control.Monad (forM_)
import Data.Text (pack)
import Data.Typeable (Typeable)
import Foreign.Ptr (Ptr)
import GHC.Records (HasField)
import Hsm.GPIO (GPIO(..), LineRequest, active, inactive, setPins)
import Hsm.Log qualified as L
import Hsm.PWM qualified as P
import System.IO.Echo (withoutInputEcho)

data Direction
  = N
  | NE
  | E
  | SE
  | S
  | SW
  | W
  | NW
  | RL
  | RR
  deriving (Read, Show, Typeable)

data Speed
  = SlowXXX
  | SlowXX
  | SlowX
  | Slow
  | Fast
  | Top
  deriving (Read, Show, Typeable)

type Duration = Float

data DriveAction
  = Move Direction Speed Duration
  | Stop Duration
  deriving (Read, Show, Typeable)

-- This function maps a `Speed` value to a corresponding PWM cycle duration.
-- It assumes a stepper motor with 200 steps per revolution, using a 1/16
-- microstep setting (so actually, 3200 steps per revolution). The returned
-- values represent the duration of a whole PWM cycle in nanoseconds.
cycleDuration :: Speed -> Int
cycleDuration SlowXXX = 8000000000 `div` 3200 -- 1/8 revs/s
cycleDuration SlowXX = 4000000000 `div` 3200 -- 1/4 revs/s
cycleDuration SlowX = 2000000000 `div` 3200 -- 1/2 revs/s
cycleDuration Slow = 1000000000 `div` 3200 -- 1 revs/s
cycleDuration Fast = 500000000 `div` 3200 -- 2 revs/s
cycleDuration Top = 250000000 `div` 3200 -- 4 revs/s

data DriveEnv = DriveEnv
  { pinEnable :: GPIO
  , pinDiag1 :: GPIO
  , pinDiag2 :: GPIO
  , pinDir1 :: GPIO
  , pinDir2 :: GPIO
  , pinDir3 :: GPIO
  , pinDir4 :: GPIO
  , pwmChannel :: P.PWMChannel
  }

driveEnvDefault :: DriveEnv
driveEnvDefault =
  DriveEnv
    { pinEnable = GPIO21
    , pinDiag1 = GPIO20
    , pinDiag2 = GPIO16
    , pinDir1 = GPIO12
    , pinDir2 = GPIO7
    , pinDir3 = GPIO8
    , pinDir4 = GPIO25
    , pwmChannel = P.PWM3
    }

drive ::
     HasField "driveEnv" env DriveEnv
  => Ptr LineRequest
  -> P.PWMHandle
  -> env
  -> [DriveAction]
  -> IO Bool
drive lineRequest pwmHandle env actions =
  withoutInputEcho $ handle handler runActions
  where
    logMsg = L.logMsg ["drive"]
    -- Sets GPIO pins to a desired state
    setCycleDuration = P.setCycleDuration pwmHandle env.driveEnv.pwmChannel
    setActive pins = setPins lineRequest (pins <*> [env.driveEnv]) active
    setInactive pins = setPins lineRequest (pins <*> [env.driveEnv]) inactive
    -- Pin assignments for various movement directions, each direction
    -- corresponds to a specific set of GPIO pins.
    pinsDiag = [pinDiag1, pinDiag2]
    pinsDir = [pinDir1, pinDir2, pinDir3, pinDir4]
    pinDiagNE = [pinDiag1]
    pinDiagSE = [pinDiag2]
    pinDiagSW = [pinDiag1]
    pinDiagNW = [pinDiag2]
    pinsN = [pinDir2, pinDir4]
    pinsE = [pinDir1, pinDir2]
    pinsS = [pinDir1, pinDir3]
    pinsW = [pinDir3, pinDir4]
    pinsRL = [pinDir1, pinDir2, pinDir3, pinDir4]
    pinsRR = []
    -- Introduces a delay with the duration converted from seconds to
    -- microseconds.
    runDelay = threadDelay . round . (* 1000000)
    -- Ensures that the system is reset to a safe state by setting the PWM
    -- cycle duration to zero and deactivating all used motor control pins.
    runRelease = do
      setCycleDuration 0
      setInactive $ pinsDir <> pinsDiag
    -- Handles each movement command and activates the appropriate pins for
    -- the requested direction. It also sets the cycle duration for PWM and
    -- holds this state for the specified duration.
    runAction (Move direction speed duration) = do
      case direction of
        N -> setActive $ pinsN <> pinsDiag
        NE -> setActive $ pinsN <> pinDiagNE
        E -> setActive $ pinsE <> pinsDiag
        SE -> setActive $ pinsS <> pinDiagSE
        S -> setActive $ pinsS <> pinsDiag
        SW -> setActive $ pinsS <> pinDiagSW
        W -> setActive $ pinsW <> pinsDiag
        NW -> setActive $ pinsN <> pinDiagNW
        RL -> setActive $ pinsRL <> pinsDiag
        RR -> setActive $ pinsRR <> pinsDiag
      setCycleDuration $ cycleDuration speed
      runDelay duration
      runRelease
    -- A Stop command causes the system to wait for the specified duration
    -- without performing any movement. During this period, the motor drivers
    -- remain enabled, effectively applying a brake to the motor by holding it
    -- in its current position.
    runAction (Stop duration) = runDelay duration
    -- Catches any asynchronous exceptions during the execution of commands.
    -- If an exception occurs, the system will log the error and ensure that
    -- the motors are safely released by calling `runRelease`.
    handler exception = do
      logMsg
        $ "Async exception caught while command was running: "
            <> pack (show @AsyncException exception)
      runRelease
      return False
    -- Executes a series of drive actions
    runActions = do
      setCycleDuration 0
      -- The A4988 motor driver is placed in sleep mode between commands to
      -- save power. To wake it up, a 1-microsecond delay is required before
      -- sending step commands. For added safety, we wait 1 millisecond to
      -- ensure the driver is fully awake.
      logMsg "Enabling drivers"
      setActive [pinEnable]
      logMsg "Allowing drivers to come out of sleep mode"
      threadDelay 1000
      forM_ actions $ \action -> do
        logMsg $ "Running action " <> pack (show action)
        runAction action
      logMsg "Disabling drivers"
      setInactive [pinEnable]
      return True