{-# LANGUAGE TypeFamilies #-}

module Hsm.Drive
  ( Drive
  , Direction (..)
  , Angle (..)
  , Speed (..)
  , Duration
  , Action (..)
  , drive
  , runDrive
  )
where

import Control.Concurrent (threadDelay)
import Control.Monad (forM_)
import Effectful (Dispatch (Static), DispatchOf, Eff, IOE, (:>))
import Effectful.Dispatch.Static (SideEffects (WithSideEffects), StaticRep, evalStaticRep, unsafeEff_)
import Effectful.Exception (AsyncException, bracket_, handle)
import Hsm.GPIO (GPIO, GPIOPin (..), active, inactive, setPins)
import Hsm.Log (Log, Severity (Attention, Info, Trace), logMsg)
import Hsm.PWM (PWM, PWMChannel (..), setCycleDuration)

data Drive (a :: * -> *) (b :: *)

type instance DispatchOf Drive = Static WithSideEffects

newtype instance StaticRep Drive
  = Drive ()

-- Defines the complete action space for omnidirectional robots:
-- - 8-directional movement (cardinal + intercardinal)
-- - Bidirectional rotation (clockwise/counter-clockwise)
data Direction
  = N
  | NE
  | E
  | SE
  | S
  | SW
  | W
  | NW
  deriving Show

data Angle
  = CCW
  | CW
  deriving Show

data Speed
  = Slow8
  | Slow4
  | Slow2
  | Slow
  | Fast
  | Top
  deriving Show

type Duration = Float

data Action
  = Move Direction Speed Duration
  | Tilt Angle Speed Duration
  | Stop Duration
  deriving Show

-- 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). Returned
-- values represent the duration of a whole PWM cycle in nanoseconds.
cycleDuration :: Speed -> Int
cycleDuration Slow8 = 8000000000 `div` 3200 -- 1/8 revs/s
cycleDuration Slow4 = 4000000000 `div` 3200 -- 1/4 revs/s
cycleDuration Slow2 = 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

ms1 :: GPIOPin
ms1 = GPIO2

ms2 :: GPIOPin
ms2 = GPIO3

ms3 :: GPIOPin
ms3 = GPIO4

notReset :: GPIOPin
notReset = GPIO27

notSleep :: GPIOPin
notSleep = GPIO22

dir1 :: GPIOPin
dir1 = GPIO10

dir2 :: GPIOPin
dir2 = GPIO9

dir3 :: GPIOPin
dir3 = GPIO11

dir4 :: GPIOPin
dir4 = GPIO5

diag2 :: GPIOPin
diag2 = GPIO6

diag1 :: GPIOPin
diag1 = GPIO13

step :: PWMChannel
step = PWM3

-- Executes a sequence of drive actions with interruption support
-- Wakes motors from SLEEP mode during execution and guarantees return to SLEEP
-- mode upon completion or interruption.
drive
  :: (GPIO :> es, IOE :> es, Log "drive" :> es, Log "gpio" :> es, Log "pwm" :> es, PWM :> es) => [Action] -> Eff es ()
drive actions =
  bracket_ awaken sleep . handle handler . forM_ actions $ \action -> do
    logMsg @"drive" Trace $ "Running action: " <> show action
    runAction action
  where
    pinsDir = [dir1, dir2, dir3, dir4]
    pinsDiag = [diag1, diag2]
    pinDiagNE = [diag1]
    pinDiagSE = [diag2]
    pinDiagSW = [diag1]
    pinDiagNW = [diag2]
    pinsN = [dir2, dir4]
    pinsE = [dir1, dir2]
    pinsS = [dir1, dir3]
    pinsW = [dir3, dir4]
    pinsCCW = [dir1, dir2, dir3, dir4]
    pinsCW = []
    -- 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.
    awaken = do
      logMsg @"drive" Trace "Enabling drivers"
      setPins [notSleep] active
      logMsg @"drive" Trace "Allowing drivers to come out of sleep mode"
      unsafeEff_ $ threadDelay 1000
    sleep = do
      logMsg @"drive" Trace "Disabling drivers"
      setPins [notSleep] inactive
    -- Catches any asynchronous exceptions during the execution of commands.
    -- If an exception occurs, the system will log the error and ensure that
    -- all pins are deactivated.
    endAction = do
      setCycleDuration step 0
      setPins (pinsDir <> pinsDiag) inactive
    handler exception = do
      logMsg @"drive" Attention $ "Async exception caught while action was running: " <> show @AsyncException exception
      endAction
    -- 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.
    delay seconds = unsafeEff_ . threadDelay . round $ seconds * 1000000
    togglePins speed duration pins = do
      setCycleDuration step $ cycleDuration speed
      setPins pins active
      delay duration
      endAction
    runAction (Move direction speed duration) =
      togglePins speed duration $
        case direction of
          N -> pinsN <> pinsDiag
          NE -> pinsN <> pinDiagNE
          E -> pinsE <> pinsDiag
          SE -> pinsS <> pinDiagSE
          S -> pinsS <> pinsDiag
          SW -> pinsS <> pinDiagSW
          W -> pinsW <> pinsDiag
          NW -> pinsN <> pinDiagNW
    runAction (Tilt angle speed duration) =
      togglePins speed duration $
        case angle of
          CCW -> pinsCCW <> pinsDiag
          CW -> pinsCW <> pinsDiag
    -- 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) = delay duration

runDrive :: (GPIO :> es, IOE :> es, Log "drive" :> es, Log "gpio" :> es) => Eff (Drive : es) a -> Eff es a
runDrive = evalStaticRep (Drive ()) . bracket_ enterActive exitActive
  where
    alwaysActive = [ms1, ms2, ms3, notReset]
    enterActive = do
      logMsg @"drive" Info "Entering active drive state"
      setPins alwaysActive active
    exitActive = do
      logMsg @"drive" Info "Exiting active drive state"
      setPins alwaysActive inactive