{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}

module RemoteControl where

import AgentTests.FunctionalAPITests (runRight)
import Control.Logger.Simple
import Crypto.Random (ChaChaDRG)
import qualified Data.Aeson as J
import qualified Data.ByteString.Char8 as B
import qualified Data.ByteString.Lazy.Char8 as LB
import Data.List (stripPrefix)
import Data.List.NonEmpty (NonEmpty (..))
import Data.Time.Clock.System (SystemTime (..))
import qualified Simplex.Messaging.Crypto as C
import Simplex.Messaging.Encoding.String (StrEncoding (..))
import Simplex.Messaging.Transport (TSbChainKeys (..))
import Simplex.Messaging.Transport.Client (TransportHost)
import qualified Simplex.RemoteControl.Client as HC (RCHostClient (action))
import qualified Simplex.RemoteControl.Client as RC
import Simplex.RemoteControl.Discovery (mkLastLocalHost, preferAddress)
import Simplex.RemoteControl.Invitation
  ( RCInvitation (..),
    RCSignedInvitation,
    verifySignedInvitation,
  )
import Simplex.RemoteControl.Types
import Test.Hspec hiding (fit, it)
import UnliftIO
import UnliftIO.Concurrent
import Util

remoteControlTests :: Spec
remoteControlTests = do
  describe "preferred bindings should go first" testPreferAddress
  describe "Invitation parsing" $ do
    it "should parse bracketed IPv6 host with port" testInvitationBracketedIPv6Host
    it "should reject bracketed non-IPv6 host" testInvitationBracketedNonIPv6HostRejected
  describe "New controller/host pairing" $ do
    it "should connect to new pairing" testNewPairing
    it "should connect to existing pairing" testExistingPairing
  describe "Multicast discovery" $ do
    it "should find paired host and connect" testMulticast

testPreferAddress :: Spec
testPreferAddress = do
  it "suppresses localhost" $
    mkLastLocalHost addrs
      `shouldBe` [ "10.20.30.40" `on` "eth0",
                   "10.20.30.42" `on` "wlan0",
                   "127.0.0.1" `on` "lo"
                 ]
  it "finds by address" $ do
    preferAddress ("127.0.0.1" `on` "lo23") addrs' `shouldBe` addrs -- localhost is back on top
    preferAddress ("10.20.30.42" `on` "wlp2s0") addrs'
      `shouldBe` [ "10.20.30.42" `on` "wlan0",
                   "10.20.30.40" `on` "eth0",
                   "127.0.0.1" `on` "lo"
                 ]
  it "finds by interface" $ do
    preferAddress ("127.1.2.3" `on` "lo") addrs' `shouldBe` addrs
    preferAddress ("0.0.0.0" `on` "eth0") addrs' `shouldBe` addrs'
  it "survives duplicates" $ do
    preferAddress ("0.0.0.0" `on` "eth1") addrsDups `shouldBe` addrsDups
    preferAddress ("0.0.0.0" `on` "eth0") ifaceDups `shouldBe` ifaceDups
  where
    on th interface = RCCtrlAddress {address = either error id $ strDecode th, interface}
    addrs =
      [ "127.0.0.1" `on` "lo", -- localhost may go first and break things
        "10.20.30.40" `on` "eth0",
        "10.20.30.42" `on` "wlan0"
      ]
    addrs' = mkLastLocalHost addrs
    addrsDups = "10.20.30.40" `on` "eth1" : addrs'
    ifaceDups = "10.20.30.41" `on` "eth0" : addrs'

testInvitationBracketedIPv6Host :: IO ()
testInvitationBracketedIPv6Host = do
  invitation <- testIPv6Invitation
  let bracketedUri =
        B.pack . replaceFirst "@2001:db8::1:" "@[2001:db8::1]:" . B.unpack $
          strEncode invitation
      expectedHost = either error id (strDecode "2001:db8::1") :: TransportHost
  case strDecode bracketedUri of
    Left err -> expectationFailure err
    Right RCInvitation {host, port} -> do
      host `shouldBe` expectedHost
      port `shouldBe` 5223

testInvitationBracketedNonIPv6HostRejected :: IO ()
testInvitationBracketedNonIPv6HostRejected = do
  invitation <- testIPv6Invitation
  let bracketedUri =
        B.pack . replaceFirst "@2001:db8::1:" "@[simplex.chat]:" . B.unpack $
          strEncode invitation
  case strDecode bracketedUri :: Either String RCInvitation of
    Left _ -> pure ()
    Right _ -> expectationFailure "expected parse failure for bracketed non-IPv6 host"

replaceFirst :: String -> String -> String -> String
replaceFirst needle replacement = go
  where
    go [] = []
    go input@(c : cs) =
      case stripPrefix needle input of
        Just rest -> replacement <> rest
        Nothing -> c : go cs

testIPv6Invitation :: IO RCInvitation
testIPv6Invitation = do
  drg <- C.newRandom
  (skey, _) <- atomically $ C.generateKeyPair @'C.Ed25519 drg
  (idkey, _) <- atomically $ C.generateKeyPair @'C.Ed25519 drg
  (dh, _) <- atomically $ C.generateKeyPair @'C.X25519 drg
  pure
    RCInvitation
      { ca = C.KeyHash "test-ca",
        host = either error id $ strDecode "2001:db8::1",
        port = 5223,
        v = supportedRCPVRange,
        app = J.String "app",
        ts = MkSystemTime 0 0,
        skey,
        idkey,
        dh
      }

testNewPairing :: IO ()
testNewPairing = do
  drg <- C.newRandom
  hp <- RC.newRCHostPairing drg
  invVar <- newEmptyMVar
  ctrlSessId <- async . runRight $ do
    logNote "c 1"
    (_found, inv, hc, r) <- RC.connectRCHost drg hp (J.String "app") False Nothing Nothing
    logNote "c 2"
    putMVar invVar (inv, hc)
    logNote "c 3"
    Right (sessId, _tls, r') <- atomically $ takeTMVar r
    logNote "c 4"
    Right (rcHostSession, _rcHelloBody, _hp') <- atomically $ takeTMVar r'
    let RCHostSession {tls, sessionKeys = HostSessKeys {chainKeys}} = rcHostSession
        TSbChainKeys {rcvKey, sndKey} = chainKeys
    sndKeyNonce <- atomically $ stateTVar sndKey C.sbcHkdf
    encCmd <- RC.rcEncryptBody sndKeyNonce "command message"
    RC.sendRCPacket tls $ LB.toStrict encCmd
    encResp <- RC.receiveRCPacket tls
    rcvKeyNonce <- atomically $ stateTVar rcvKey C.sbcHkdf
    resp <- RC.rcDecryptBody rcvKeyNonce $ LB.fromStrict encResp
    liftIO $ resp `shouldBe` "response message"
    logNote "c 5"
    threadDelay 250000
    logNote "ctrl: ciao"
    liftIO $ RC.cancelHostClient hc
    pure sessId

  (signedInv, hc) <- takeMVar invVar
  -- logNote $ decodeUtf8 $ strEncode inv
  inv <- maybe (fail "bad invite") pure $ verifySignedInvitation signedInv

  hostSessId <- async . runRight $ do
    logNote "h 1"
    (rcCtrlClient, r) <- RC.connectRCCtrl drg inv Nothing (J.String "app")
    logNote "h 2"
    Right (sessId', _tls, r') <- atomically $ takeTMVar r
    logNote "h 3"
    liftIO $ RC.confirmCtrlSession rcCtrlClient True
    logNote "h 4"
    Right (rcCtrlSession, _rcCtrlPairing) <- atomically $ takeTMVar r'
    let RCCtrlSession {tls, sessionKeys = CtrlSessKeys {chainKeys}} = rcCtrlSession
        TSbChainKeys {rcvKey, sndKey} = chainKeys
    encCmd <- RC.receiveRCPacket tls
    rcvKeyNonce <- atomically $ stateTVar rcvKey C.sbcHkdf
    cmd <- RC.rcDecryptBody rcvKeyNonce $ LB.fromStrict encCmd
    liftIO $ cmd `shouldBe` "command message"
    sndKeyNonce <- atomically $ stateTVar sndKey C.sbcHkdf
    encResp <- RC.rcEncryptBody sndKeyNonce "response message"
    RC.sendRCPacket tls $ LB.toStrict encResp
    logNote "h 5"
    threadDelay 250000
    logNote "ctrl: adios"
    pure sessId'

  waitCatch (HC.action hc) >>= \case
    Left err -> fromException err `shouldBe` Just AsyncCancelled
    Right () -> fail "Unexpected controller finish"

  timeout 5000000 (waitBoth ctrlSessId hostSessId) >>= \case
    Just (sessId, sessId') -> sessId `shouldBe` sessId'
    _ -> fail "timeout"

testExistingPairing :: IO ()
testExistingPairing = do
  drg <- C.newRandom
  invVar <- newEmptyMVar
  hp <- RC.newRCHostPairing drg
  ctrl <- runCtrl drg False hp invVar
  inv <- takeMVar invVar
  let cp_ = Nothing
  host <- runHostURI drg cp_ inv
  timeout 5000000 (waitBoth ctrl host) >>= \case
    Nothing -> fail "timeout"
    Just (hp', cp') -> do
      ctrl' <- runCtrl drg False hp' invVar
      inv' <- takeMVar invVar
      host' <- runHostURI drg (Just cp') inv'
      timeout 5000000 (waitBoth ctrl' host') >>= \case
        Nothing -> fail "timeout"
        Just (_hp2, cp2) -> do
          ctrl2 <- runCtrl drg False hp' invVar -- old host pairing used to test controller not updating state
          inv2 <- takeMVar invVar
          host2 <- runHostURI drg (Just cp2) inv2
          timeout 5000000 (waitBoth ctrl2 host2) >>= \case
            Nothing -> fail "timeout"
            Just (hp3, cp3) -> do
              ctrl3 <- runCtrl drg False hp3 invVar
              inv3 <- takeMVar invVar
              host3 <- runHostURI drg (Just cp3) inv3
              timeout 5000000 (waitBoth ctrl3 host3) >>= \case
                Nothing -> fail "timeout"
                Just _ -> pure ()

testMulticast :: IO ()
testMulticast = do
  drg <- C.newRandom
  subscribers <- newTMVarIO 0
  invVar <- newEmptyMVar
  hp <- RC.newRCHostPairing drg
  ctrl <- runCtrl drg False hp invVar
  inv <- takeMVar invVar
  let cp_ = Nothing
  host <- runHostURI drg cp_ inv
  timeout 5000000 (waitBoth ctrl host) >>= \case
    Nothing -> fail "timeout"
    Just (hp', cp') -> do
      ctrl' <- runCtrl drg True hp' invVar
      _inv <- takeMVar invVar
      host' <- runHostMulticast drg subscribers cp'
      timeout 5000000 (waitBoth ctrl' host') >>= \case
        Nothing -> fail "timeout"
        Just _ -> pure ()

runCtrl :: TVar ChaChaDRG -> Bool -> RCHostPairing -> MVar RCSignedInvitation -> IO (Async RCHostPairing)
runCtrl drg multicast hp invVar = async . runRight $ do
  (_found, inv, hc, r) <- RC.connectRCHost drg hp (J.String "app") multicast Nothing Nothing
  putMVar invVar inv
  Right (_sessId, _tls, r') <- atomically $ takeTMVar r
  Right (_rcHostSession, _rcHelloBody, hp') <- atomically $ takeTMVar r'
  threadDelay 250000
  liftIO $ RC.cancelHostClient hc
  pure hp'

runHostURI :: TVar ChaChaDRG -> Maybe RCCtrlPairing -> RCSignedInvitation -> IO (Async RCCtrlPairing)
runHostURI drg cp_ signedInv = async . runRight $ do
  inv <- maybe (fail "bad invite") pure $ verifySignedInvitation signedInv
  (rcCtrlClient, r) <- RC.connectRCCtrl drg inv cp_ (J.String "app")
  Right (_sessId', _tls, r') <- atomically $ takeTMVar r
  liftIO $ RC.confirmCtrlSession rcCtrlClient True
  Right (_rcCtrlSession, cp') <- atomically $ takeTMVar r'
  threadDelay 250000
  pure cp'

runHostMulticast :: TVar ChaChaDRG -> TMVar Int -> RCCtrlPairing -> IO (Async RCCtrlPairing)
runHostMulticast drg subscribers cp = async . runRight $ do
  (pairing, inv) <- RC.discoverRCCtrl subscribers (cp :| [])
  (rcCtrlClient, r) <- RC.connectRCCtrl drg inv (Just pairing) (J.String "app")
  Right (_sessId', _tls, r') <- atomically $ takeTMVar r
  liftIO $ RC.confirmCtrlSession rcCtrlClient True
  Right (_rcCtrlSession, cp') <- atomically $ takeTMVar r'
  threadDelay 250000
  pure cp'