{-# LANGUAGE TemplateHaskell #-}

module Literate.Web.Live.WebSocket where

import Control.Applicative (Alternative (..), Applicative (..))
import Control.Concurrent.Async qualified as Async
import Control.Concurrent.STM (STM)
import Control.Concurrent.STM qualified as STM
import Control.Exception qualified as Exception
import Control.Monad (Monad (..))
import Data.ByteString.Lazy qualified as BSL
import Data.Either (Either (..))
import Data.FileEmbed (embedFile)
import Data.Function (($))
import Data.Functor ((<$>))
import Data.Map.Strict qualified as Map
import Data.Maybe (Maybe (..))
import Data.Monoid (mconcat)
import Data.Semigroup (Semigroup (..))
import Data.Text (Text)
import Data.Text.Lazy qualified as TextL
import Data.Text.Lazy.Encoding qualified as TextL
import Data.Text.Lazy.IO qualified as TextL
import Network.WebSockets (ConnectionException)
import Network.WebSockets qualified as WS
import System.IO (IO)
import System.IO qualified as Sys
import Text.Show (show)

import Literate.Web qualified as Web
import Literate.Web.Live.Asset
import Literate.Web.Live.Common

webSocketApp ::
  model ->
  STM model ->
  (model -> IO (Map.Map Web.OutputPath (IO BSL.ByteString))) ->
  WS.ServerApp
webSocketApp iniModel modelSTM siteMap pendingConn =
  withWebSocket pendingConn $ \(routeSTM, routeSend) -> do
    let loop model route = do
          STM.atomically (Left <$> modelSTM <|> Right <$> routeSTM) >>= \case
            Left newModel -> loop newModel route
            Right newRoute -> loop model newRoute
          site <- siteMap model
          TextL.hPutStrLn Sys.stderr $ "webSocketApp: route: " <> TextL.pack (show route)
          case Map.lookup route site of
            Nothing -> routeSend [WS.Binary (liveErrorHtmlResponse decodeRouteNothingMsg)]
            Just contentIO -> do
              content <- contentIO
              routeSend [WS.Binary content]
    loop
      iniModel
      Web.OutputPath
        { Web.outputPathSegs = ["index"]
        , Web.outputPathExts = ["html"]
        }

withWebSocket ::
  WS.PendingConnection ->
  ((STM Web.OutputPath, [WS.DataMessage] -> IO ()) -> IO ()) ->
  IO ()
withWebSocket pendingConn f = do
  var <- STM.newEmptyTMVarIO
  conn :: WS.Connection <- WS.acceptRequest pendingConn
  WS.withPingThread conn 30 (pure ()) $
    Async.withAsync
      ( let loop =
              Exception.try @ConnectionException (WS.receiveData @Text conn) >>= \case
                Left (WS.CloseRequest _ (TextL.decodeUtf8 -> reason)) -> do
                  TextL.hPutStrLn Sys.stderr $ "WebSocket.CloseRequest, reason: " <> reason
                Left WS.ConnectionClosed -> do
                  Sys.hPrint Sys.stderr (["withWebSocket", "loop", "ConnectionClosed"] :: [Text])
                Left WS.ParseException{} -> do
                  TextL.hPutStrLn Sys.stderr $ "WebSocket.ParseException"
                Left WS.UnicodeException{} -> do
                  TextL.hPutStrLn Sys.stderr $ "WebSocket.UnicodeException"
                Right msg -> do
                  STM.atomically $ STM.writeTMVar var $ decodeOutputPath msg
                  loop
        in loop
      )
      Async.wait
  Exception.try @ConnectionException
    (f (STM.readTMVar var, WS.sendDataMessages conn))
    >>= \case
      Left (WS.CloseRequest _ (TextL.decodeUtf8 -> reason)) -> do
        TextL.hPutStrLn Sys.stderr $ "WebSocket.CloseRequest, reason: " <> reason
      Left WS.ConnectionClosed -> do
        Sys.hPrint Sys.stderr (["withWebSocket", "f", "ConnectionClosed"] :: [Text])
      Left WS.ParseException{} -> do
        TextL.hPutStrLn Sys.stderr $ "WebSocket.ParseException"
      Left WS.UnicodeException{} -> do
        TextL.hPutStrLn Sys.stderr $ "WebSocket.UnicodeException"
      Right () -> WS.sendClose conn ("Bye!" :: Text)

-- case connExc of
--  WS.CloseRequest _ (Text.decodeUtf8 -> reason) ->
--    log LevelInfo $ "Closing websocket connection (reason: " <> reason <> ")"
--  _ ->
--    log LevelError $ "Websocket error: " <> show connExc
-- LVar.removeListener model subId

-- | A handler takes a websocket connection and the current model and then watches
--    for websocket messages. It must return a new route to watch (after that, the
--    returned route's HTML will be sent back to the client).
--
--   Note that this is usually a long-running thread that waits for the client's
--   messages. But you can also use it to implement custom server actions, by handling
--   the incoming websocket messages or other IO events in any way you like.
--
--   Also note that whenever the model is updated, the handler action will be
--   stopped and then restarted with the new model as argument.
newtype WebSocketHandler model = WebSocketHandler
  { unWebSocketHandler :: WS.Connection {-RO IO-} -> model -> IO Text
  }

webSocketHandler = WebSocketHandler $ \conn _model -> do
  msg :: Text <- WS.receiveData conn
  -- log LevelDebug $ "<~~ " <> show msg
  pure msg

-- where log lvl (t :: Text) = logWithoutLoc "ema.ws" lvl t

-- Browser-side JavaScript code for interacting with the Haskell server
webSocketClientJS :: BSL.ByteString
webSocketClientJS =
  mconcat
    [ TextL.encodeUtf8
        ( mconcat
            [ "<script type='module' src='https://cdn.jsdelivr.net/npm/morphdom@2.7.2/dist/morphdom-umd.min.js'></script>"
            , "<script type='module'>"
            ]
        )
    , BSL.fromStrict $(embedFile "www/live-shim.js")
    , TextL.encodeUtf8
        ( mconcat
            [ "window.onpageshow = function () { init(false) };"
            , "</script>"
            ]
        )
    ]

-- instance WS.WebSocketsData Web.OutputPath where

{-
    --subId <- LVar.addListener model
    -- let log lvl (s :: Text) = logWithoutLoc (toText @String $ printf "ema.ws.%.2d" subId) lvl s
    -- log LevelInfo "Connected"
    let
      wsHandler = unWebSocketHandler webSockH conn
      sendRouteHtmlToClient path s = do
        case Map.lookup Web.siteMap of
          Nothing -> routeRW.write $ liveErrorHtmlResponse "NoSuchRoute"
        {-
        decodeUrlRoute @model s path & \case
          Left err -> do
            -- log LevelError $ badRouteEncodingMsg err
          Right Nothing ->
            MC.exec @IO $ WS.sendTextData conn $ liveErrorHtmlResponse decodeRouteNothingMsg
          Right (Just r) -> do
            renderCatchingErrors s r >>= \case
              AssetGenerated Html html ->
                MC.exec @IO $ WS.sendTextData conn $ html <> toLazy wsClientHtml
              -- HACK: We expect the websocket client should check for REDIRECT prefix.
              -- Not bothering with JSON response to avoid having to JSON parse every HTML dump.
              AssetStatic _staticPath ->
                MC.exec @IO $ WS.sendTextData conn $ "REDIRECT " <> toText (review (fromPrism_ $ routePrism s) r)
              AssetGenerated Other _s ->
                MC.exec @IO $ WS.sendTextData conn $ "REDIRECT " <> toText (review (fromPrism_ $ routePrism s) r)
        -}
      -- log LevelDebug $ " ~~> " <> show r
      -- @mWatchingRoute@ is the route currently being watched.
      loop mWatchingRoute = do
        -- Listen *until* either we get a new value, or the client requests
        -- to switch to a new route.
        wsHandler currentModel >>= \mNextRoute ->
          -- The user clicked on a route link; send them the HTML for
          -- that route this time, ignoring what we are watching
          -- currently (we expect the user to initiate a watch route
          -- request immediately following this).
          sendRouteHtmlToClient mNextRoute ()
          loop mNextRoute
        {-
        currentModel <- LVar.get model
        race (LVar.listenNext model subId) (wsHandler currentModel) >>= \case
          Left newModel -> do
            -- The page the user is currently viewing has changed. Send
            -- the new HTML to them.
            sendRouteHtmlToClient mWatchingRoute newModel
            loop mWatchingRoute
          Right mNextRoute -> do
            -- The user clicked on a route link; send them the HTML for
            -- that route this time, ignoring what we are watching
            -- currently (we expect the user to initiate a watch route
            -- request immediately following this).
            sendRouteHtmlToClient mNextRoute =<< LVar.get model
            loop mNextRoute
        -}
    -- Wait for the client to send the first request with the initial route.
    --modelRO.React.onWrite $ \model -> do
    --  sendRouteHtmlToClient <$> routeRO.React.read <*> pure model
    --routeRO.React.onWrite $ \route ->
    --  sendRouteHtmlToClient <$> route <*> modelRO.React.read
    --  Exception.try @ConnectionException (loop mInitialRoute) >>= \case
    --    Right () -> pass
    --    Left _connExc -> do
    --      return ()
choose :: [(STM a, a -> IO ())] -> IO ()
choose choices = do
    act <- atomically (foldr1 orElse actions)
    act
  where
    actions :: [STM (IO ())]
    actions = [ do val <- guard
                   return (rhs val)
              | (guard, rhs) <- choices ]
atomicallyStubborn :: forall a. STM a -> IO a
atomicallyStubborn stm = persevere
  where
    persevere :: IO a
    persevere =
        catch (atomically stm) $ \BlockedIndefinitelyOnSTM ->
          persevere
dequeue :: TQueue (Either NetworkFailure Int) -> IO (Either NetworkFailure Int)
dequeue queue = do
    -- consider the write thread to be a GC root when it is blocked on the TQueue; this way it will not be considered as stalled, and nor will any other threads that depend on it (such as the fourth thread in the example above). We can do this by providing creating a stable pointer to the thread (this workaround is due to Simon Marlow)
    tid <- myThreadId
    bracket (newStablePtr tid) freeStablePtr $ \_ ->
      atomically $ readTQueue queue
-}