home-assistant-controller/src/HomeAssistant/Runtime/Connection.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

module HomeAssistant.Runtime.Connection
  ( readerAction
  , writerAction
  , encodeService
  , dedupeBatch
  ) where

import Control.Concurrent.STM
  ( TChan
  , atomically
  , readTChan
  , readTVar
  , retry
  , tryReadTChan
  , writeTChan
  , writeTVar
  )
import Control.Concurrent.Async (race)
import Control.Concurrent (threadDelay)
import Control.Exception (onException)
import Control.Exception.Annotated (throw)
import Control.Lens ((^?))
import Control.Monad (forever, forM_)
import Data.Aeson (Value, eitherDecode, encode, object, (.=))
import Data.Aeson.Lens (key, _String)
import Data.List (sort)
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import qualified Data.Text as T
import Data.Void (Void)
import HomeAssistant.Controller (Service (..), Target (..))
import HomeAssistant.Runtime.Bus
import HomeAssistant.Runtime.Supervisor (Fatal (..))
import qualified Network.WebSockets as WS
import Katip (runKatipContextT, sl, logFM, Severity (..), ls)
import Data.UUID (toText)
import AFRP (Request(..), Event(..))

-- | Connect, authenticate, subscribe, then receive and broadcast forever.
-- Restarting this action reconnects. All setup sends happen before the
-- connection is published in the bus, so only the writer sends afterwards.
readerAction :: String -> Int -> String -> S.Set T.Text -> Bus -> IO Void
readerAction host port token ents bus =
  WS.runClient host port "/api/websocket" $ \conn -> do
    handshake conn token
    subscribe bus conn ents
    atomically $ writeTVar (busConn bus) (Just conn)
    putStrLn "[reader] connected"
    -- Unpublish on exit so the writer blocks and the backlog survives the outage.
    receiveLoop bus conn `onException` atomically (writeTVar (busConn bus) Nothing)

handshake :: WS.Connection -> String -> IO ()
handshake conn token = do
  required <- receiveJSON conn
  expectType "auth_required" required
  WS.sendTextData conn $ encode $ object
    [ "type"         .= ("auth" :: T.Text)
    , "access_token" .= token
    ]
  ok <- receiveJSON conn
  expectType "auth_ok" ok

expectType :: T.Text -> Value -> IO ()
expectType expected msg =
  case msg ^? key "type" . _String of
    Just t | t == expected -> pure ()
    _ -> throw (Fatal $ "expected " <> expected <> ", got: " <> T.pack (show msg))

subscribe :: Bus -> WS.Connection -> S.Set T.Text -> IO ()
subscribe bus conn ents =
  forM_ (S.toList ents) $ \entityId -> do
    print entityId
    sid <- generateCallId (busGen bus)
    WS.sendTextData conn $ encode $ object
      [ "id"         .= sid
      , "type"       .= ("subscribe_trigger" :: T.Text)
      , "trigger"    .= object
        [ "platform" .= ("state" :: T.Text)
        , "entity_id" .= entityId
        ]
      ]

-- | Undecodable messages are skipped: reconnecting cannot fix a decode
-- problem, so crashing here would only produce a hot restart loop.
--
-- Each read races a one-second timeout: a timeout broadcasts 'Tick' so
-- time-based primitives (debounce, rollup, fixed, ...) keep advancing
-- even when no state changes arrive.
receiveLoop :: Bus -> WS.Connection -> IO Void
receiveLoop bus conn = forever $ do
  winner <- race (threadDelay 1000000) (WS.receiveData conn)
  case winner of
    Left () -> atomically $ writeTChan (busInbound bus) Tick
    Right msg -> case eitherDecode msg of
      Left err -> putStrLn $ "[reader] skipping undecodable message: " <> err
      Right v  -> atomically $ writeTChan (busInbound bus) (Event v)

receiveJSON :: WS.Connection -> IO Value
receiveJSON conn = do
  msg <- WS.receiveData conn
  case eitherDecode msg of
    Left err -> throw (Fatal $ "Invalid JSON from Home Assistant: " <> T.pack err)
    Right x  -> pure x

-- | Floor between sends within a batch: 100ms, so a many-distinct-target
-- flood still caps at ~10 sends/sec even after dedupe.
minInterval :: Int
minInterval = 100000

-- | Non-blocking drain of everything queued on a channel. Returns items
-- oldest-first (FIFO from the channel), so prepending the blocking
-- `readTChan` item keeps the whole batch oldest-first for `dedupeBatch`.
drainTry :: TChan a -> IO [a]
drainTry chan = go []
  where
    go acc = do
      m <- atomically $ tryReadTChan chan
      case m of
        Nothing -> pure (reverse acc)
        Just x  -> go (x : acc)

sendWithId :: Bus -> WS.Connection -> Request -> Service -> IO ()
sendWithId bus conn request svc = do
  callId <- generateCallId (busGen bus)
  let textData = encode $ encodeService callId svc
  runKatipContextT (busLogEnv bus) (sl "traceId" (toText (requestTraceId request))) "connection" $
    logFM DebugS (ls textData)
  WS.sendTextData conn textData

writerAction :: Bus -> IO Void
writerAction bus = forever $ do
  first <- atomically $ readTChan (busOutbound bus)
  rest  <- drainTry (busOutbound bus)
  let deduped = dedupeBatch (first : rest)
  conn <- atomically $ readTVar (busConn bus) >>= maybe retry pure
  forM_ deduped $ \(request, svc) -> do
    sendWithId bus conn request svc
    threadDelay minInterval

encodeService :: Int -> Service -> Value
encodeService callId Service{..} = object $
  [ "id"      .= callId
  , "type"    .= ("call_service" :: T.Text)
  , "domain"  .= serviceDomain
  , "service" .= serviceName
  , "target"  .= targetObject serviceTarget
  ] <> maybe [] (\d -> ["service_data" .= d]) serviceData

-- | Collapse a drained batch of outbound calls: the newest call per
-- `(domain, service, sorted-targets)` survives; older duplicates are
-- dropped. `serviceData` is not part of the key, so a newer `turn_on`
-- with different brightness supersedes an older one to the same target.
dedupeBatch :: [(Request, Service)] -> [(Request, Service)]
dedupeBatch = M.elems . foldl' ins M.empty
  where
    ins m (req, svc) = M.insert (dedupeKey svc) (req, svc) m

dedupeKey :: Service -> (T.Text, T.Text, [Target])
dedupeKey Service{..} = (serviceDomain, serviceName, sort serviceTarget)

-- | A single target encodes as a scalar; multiple encode as a list. Empty
-- lists are omitted so Home Assistant receives only populated keys.
targetObject :: [Target] -> Value
targetObject targets = object $ entityPart <> areaPart
  where
    entityPart = emit "entity_id" [e | EntityId e <- targets]
    areaPart   = emit "area_id" [a | AreaId a <- targets]
    emit _ []     = []
    emit k [x]    = [k .= x]
    emit k many   = [k .= many]