{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}

module Hsm.Battery
  ( Battery
  , BatteryState (..)
  , readBatteryState
  , runBattery
  )
where

import Control.Exception (throwIO)
import Control.Monad (when)
import Data.Bits ((.&.), (.>>.), (.|.))
import Data.Vector.Storable (fromList, unsafeWith)
import Data.Word (Word16, Word8, byteSwap16)
import Effectful (Dispatch (Static), DispatchOf, Eff, IOE, liftIO, (:>))
import Effectful.Dispatch.Static (SideEffects (WithSideEffects), StaticRep, evalStaticRep, getStaticRep, unsafeEff_)
import Effectful.Exception (Exception, bracket)
import Foreign.Marshal.Utils (new, with)
import Foreign.Ptr (castPtr)
import Foreign.Storable (peek)
import Hsm.Battery.FFI (I2CMsg (..), I2CRdWrIoctlData (..), i2cMNoStart, i2cMRd, i2cRdWr, ioctl)
import Hsm.Core.Show (showHex)
import Hsm.Log (Log, Severity (Info, Trace), logMsg)
import System.Posix.IO (OpenMode (ReadWrite), closeFd, defaultFileFlags, openFd)
import System.Posix.Types (Fd)

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

type instance DispatchOf Battery = Static WithSideEffects

newtype instance StaticRep Battery
  = Battery Fd

data BatteryState = BatteryState
  { voltage :: Float
  , current :: Float
  , power :: Float
  }
  deriving Show

data BatteryException
  = IOCtlErrorReturnValue
  deriving (Exception, Show)

ina226Addr :: Word16
ina226Addr = 0x40

i2cWriteShort :: (Battery :> es, Log "battery" :> es) => Word8 -> Word16 -> Eff es ()
i2cWriteShort command bytes = do
  Battery fd <- getStaticRep
  logMsg Trace $ "Sending I2C write command " <> showHex command <> " with data " <> showHex bytes
  unsafeEff_ . unsafeWith commandVec $ \commandPtr ->
    with (msgVec commandPtr) $ \msgPtr ->
      with (ioVec msgPtr) $ \ioPtr -> do
        res <- ioctl fd i2cRdWr ioPtr
        when (res == (-1)) $ throwIO IOCtlErrorReturnValue
  where
    commandVec = fromList [command, fromIntegral $ (bytes .&. 0xff00) .>>. 8, fromIntegral $ bytes .&. 0xff]
    msgVec = I2CMsg ina226Addr 0 3
    ioVec msgPtr = I2CRdWrIoctlData msgPtr 1

i2cReadShort :: (Battery :> es, Log "battery" :> es) => Word8 -> Eff es Word16
i2cReadShort command = do
  Battery fd <- getStaticRep
  logMsg Trace $ "Sending I2C read command " <> showHex command
  unsafeEff_ $ do
    resPtr <- new 0
    with command $ \commandPtr ->
      unsafeWith (msgVec commandPtr resPtr) $ \msgPtr ->
        with (ioVec msgPtr) $ \ioPtr -> do
          res <- ioctl fd i2cRdWr ioPtr
          when (res == (-1)) $ throwIO IOCtlErrorReturnValue
          byteSwap16 <$> peek resPtr
  where
    i2cFlags = fromIntegral $ i2cMRd .|. i2cMNoStart
    msgVec commandPtr resPtr = fromList [I2CMsg ina226Addr 0 1 commandPtr, I2CMsg ina226Addr i2cFlags 2 $ castPtr resPtr]
    ioVec msgPtr = I2CRdWrIoctlData msgPtr 2

readBatteryState :: (Battery :> es, Log "battery" :> es) => Eff es BatteryState
readBatteryState = do
  voltage <- (* 0.00125) . fromIntegral <$> i2cReadShort ina226BusVoltageReg
  current <- (* 0.0005) . fromIntegral <$> i2cReadShort ina226CurrentReg
  power <- (* 0.0125) . fromIntegral <$> i2cReadShort ina226PowerReg
  return BatteryState{..}
  where
    ina226BusVoltageReg = 0x02
    ina226PowerReg = 0x03
    ina226CurrentReg = 0x04

runBattery :: (IOE :> es, Log "battery" :> es) => Eff (Battery : es) a -> Eff es a
runBattery action =
  bracket openI2CDev closeI2CDev $ \fd ->
    evalStaticRep (Battery fd) $ do
      i2cWriteShort ina226CfgReg 0x8000 -- reset
      i2cWriteShort ina226CfgReg 0x4527 -- average over 16 samples
      i2cWriteShort ina226CalReg 1024 -- 1A, 0.1 Ohm resistor
      manId <- i2cReadShort ina226ManufacturerId
      dieId <- i2cReadShort ina226DieId
      cfgReg <- i2cReadShort ina226CfgReg
      logMsg Info $ "INA226 manufacturer ID: " <> showHex manId
      logMsg Info $ "INA226 die ID: " <> showHex dieId
      logMsg Info $ "INA226 configuration register: " <> showHex cfgReg
      action
  where
    devPath = "/dev/i2c-0"
    openI2CDev = do
      logMsg Info $ "Opening I2C INA226 device at " <> devPath
      fd <- liftIO $ openFd devPath ReadWrite defaultFileFlags
      logMsg Info $ "I2C INA226 device opened with FD: " <> show fd
      return fd
    closeI2CDev fd = do
      logMsg Info $ "Closing I2C INA226 device at " <> devPath <> " with FD: " <> show fd
      liftIO $ closeFd fd
    ina226CfgReg = 0x00
    ina226CalReg = 0x05
    ina226ManufacturerId = 0xfe
    ina226DieId = 0xff