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

import Control.Concurrent
import Control.Monad.IO.Class
import Data.Maybe
import Data.Text (Text)
import Data.UUID
import Network.Consul (deleteKey,getKey,getSessionInfo,initializeConsulClient,putKey,withSession,ConsulClient(..),runService,getServiceHealth)
import Network.Consul.Types
import qualified Network.Consul.Internal as I
import Network.HTTP.Client
import Network.Socket (PortNumber(..))
import System.Random
import Test.Tasty
import Test.Tasty.HUnit

client :: IO ConsulClient
client = initializeConsulClient "localhost" 8500 Nothing

{- Internal Tests -}
internalKVTests :: TestTree
internalKVTests = testGroup "Internal Key Value" [testGetInvalidKey, testPutKey,
  testGetKey,testGetKeys,testListKeys,testDeleteKey,testGetNullValueKey,testDeleteRecursive]

testGetInvalidKey :: TestTree
testGetInvalidKey = testCase "testGetInvalidKey" $ do
  _client@ConsulClient{..} <- client
  x <- I.getKey ccManager (I.hostWithScheme _client) ccPort "nokey" Nothing Nothing Nothing
  assertEqual "testGetInvalidKey: Found a key that doesn't exist" x Nothing

testPutKey :: TestTree
testPutKey = testCase "testPutKey" $ do
  _client@ConsulClient{..} <- client
  let put = KeyValuePut "/testPutKey" "Test" Nothing Nothing
  x <- I.putKey ccManager (I.hostWithScheme _client) ccPort put Nothing
  assertEqual "testPutKey: Write failed" True x

testGetKey :: TestTree
testGetKey = testCase "testGetKey" $ do
  _client@ConsulClient{..} <- client
  let put = KeyValuePut "/testGetKey" "Test" Nothing Nothing
  x1 <- I.putKey ccManager (I.hostWithScheme _client) ccPort put Nothing
  assertEqual "testGetKey: Write failed" True x1
  x2 <- I.getKey ccManager (I.hostWithScheme _client) ccPort "/testGetKey" Nothing Nothing Nothing
  case x2 of
    Just x -> assertEqual "testGetKey: Incorrect Value" (kvValue x) (Just "Test")
    Nothing -> assertFailure "testGetKey: No value returned"

testGetNullValueKey :: TestTree
testGetNullValueKey = testCase "testGetNullValueKey" $ do
  _client@ConsulClient{..} <- client
  let put = KeyValuePut "/testGetNullValueKey" "" Nothing Nothing
  x1 <- I.putKey ccManager (I.hostWithScheme _client) ccPort put Nothing
  assertEqual "testGetNullValueKey: Write failed" True x1
  liftIO $ threadDelay (500 * 1000)
  x2 <- I.getKey ccManager (I.hostWithScheme _client) ccPort "/testGetNullValueKey" Nothing Nothing Nothing
  case x2 of
    Just x -> assertEqual "testGetNullValueKey: Incorrect Value" (kvValue x) Nothing
    Nothing -> assertFailure "testGetNullValueKey: No value returned"

testGetKeys :: TestTree
testGetKeys = testCase "testGetKeys" $ do
  _client@ConsulClient{..} <- client
  let put1 = KeyValuePut "/testGetKeys/key1" "Test" Nothing Nothing
  x1 <- I.putKey ccManager (I.hostWithScheme _client) ccPort put1 Nothing
  assertEqual "testGetKeys: Write failed" True x1
  let put2 = KeyValuePut "/testGetKeys/key2" "Test" Nothing Nothing
  x2 <- I.putKey ccManager (I.hostWithScheme _client) ccPort put2 Nothing
  assertEqual "testGetKeys: Write failed" True x2
  x3 <- I.getKeys ccManager (I.hostWithScheme _client) ccPort "/testGetKeys" Nothing Nothing Nothing
  assertEqual "testGetKeys: Incorrect number of results" 2 (length x3)

testListKeys :: TestTree
testListKeys = testCase "testListKeys" $ do
  _client@ConsulClient{..} <- client
  let put1 = KeyValuePut "/testListKeys/key1" "Test" Nothing Nothing
  x1 <- I.putKey ccManager (I.hostWithScheme _client) ccPort put1 Nothing
  assertEqual "testListKeys: Write failed" True x1
  let put2 = KeyValuePut "/testListKeys/key2" "Test" Nothing Nothing
  x2 <- I.putKey ccManager (I.hostWithScheme _client) ccPort put2 Nothing
  assertEqual "testListKeys: Write failed" True x2
  x3 <- I.listKeys ccManager (I.hostWithScheme _client) ccPort "/testListKeys/" Nothing Nothing Nothing
  assertEqual "testListKeys: Incorrect number of results" 2 (length x3)

testDeleteKey :: TestTree
testDeleteKey = testCase "testDeleteKey" $ do
  _client@ConsulClient{..} <- client
  let put1 = KeyValuePut "/testDeleteKey" "Test" Nothing Nothing
  x1 <- I.putKey ccManager (I.hostWithScheme _client) ccPort put1 Nothing
  assertEqual "testDeleteKey: Write failed" True x1
  x2 <- I.deleteKey ccManager (I.hostWithScheme _client) ccPort "/testDeleteKey" False Nothing
  assertEqual "testDeleteKey: Delete Failed" True x2
  x3 <- I.getKey ccManager (I.hostWithScheme _client) ccPort "/testDeleteKey" Nothing Nothing Nothing
  assertEqual "testDeleteKey: Key was not deleted" Nothing x3

testDeleteRecursive :: TestTree
testDeleteRecursive = testCase "testDeleteRecursive" $ do
  _client@ConsulClient{..} <- client
  let put1 = KeyValuePut "/testDeleteRecursive/1" "Test" Nothing Nothing
      put2 = KeyValuePut "/testDeleteRecursive/2" "Test" Nothing Nothing
  x1 <- I.putKey ccManager (I.hostWithScheme _client) ccPort put1 Nothing
  assertEqual "testDeleteKey: Write failed" True x1
  x2 <- I.putKey ccManager (I.hostWithScheme _client) ccPort put2 Nothing
  assertEqual "testDeleteKey: Write failed" True x2
  I.deleteKey ccManager (I.hostWithScheme _client) ccPort "/testDeleteRecursive/" True Nothing
  x3 <- I.getKey ccManager (I.hostWithScheme _client) ccPort "/testDeleteRecursive/1" Nothing Nothing Nothing
  assertEqual "testDeleteKey: Key was not deleted" Nothing x3

{- Client KV -}
clientKVTests :: TestTree
clientKVTests = testGroup "Client KV Tests" [testDeleteRecursiveClient]

testDeleteRecursiveClient :: TestTree
testDeleteRecursiveClient = testCase "testDeleteRecursiveClient" $ do
  c <- client
  let put1 = KeyValuePut "/testDeleteRecursive/1" "Test" Nothing Nothing
      put2 = KeyValuePut "/testDeleteRecursive/2" "Test" Nothing Nothing
  x1 <- putKey c put1 Nothing
  assertEqual "testDeleteKey: Write failed" True x1
  x2 <- putKey c put2 Nothing
  assertEqual "testDeleteKey: Write failed" True x2
  deleteKey c "/testDeleteRecursive/" True Nothing
  x3 <- getKey c "/testDeleteRecursive/1" Nothing Nothing Nothing
  assertEqual "testDeleteKey: Key was not deleted" Nothing x3

{- Agent -}
testRegisterService :: TestTree
testRegisterService = testCase "testRegisterService" $ do
  _client@ConsulClient{..} <- client
  let req = RegisterService Nothing "testService" ["test"] Nothing (Just $ Ttl "10s")
  val <- I.registerService ccManager (I.hostWithScheme _client) ccPort req Nothing
  assertEqual "testRegisterService: Service was not created" val True
  mService <- I.getService ccManager (I.hostWithScheme _client) ccPort "testService" Nothing Nothing
  case mService of
    Just _ -> return ()
    Nothing -> assertFailure "testRegisterService: Service was not found"

testGetSelf :: TestTree
testGetSelf = testCase "testGetSelf" $ do
  _client@ConsulClient{..} <- client
  x <- I.getSelf ccManager (I.hostWithScheme _client) ccPort
  assertEqual "testGetSelf: Self not returned" True (isJust x)

{-
testRegisterHealthCheck :: TestTree
testRegisterHealthCheck = testCase "testRegisterHealthCheck" $ do
  _client@ConsulClient{..} <- client
  let check = RegisterHealthCheck "testHealthCheck" "testHealthCheck" "" Nothing Nothing (Just "15s")
  x1 <- I.registerHealthCheck ccManager (I.hostWithScheme _client) ccPort check
  undefined -}

{- Health Checks -}
testGetServiceHealth :: TestTree
testGetServiceHealth = testCase "testGetServiceHealth" $ do
  _client@ConsulClient{..} <- client
  let req = RegisterService (Just "testGetServiceHealth") "testGetServiceHealth" [] Nothing Nothing
  r1 <- I.registerService ccManager (I.hostWithScheme _client) ccPort req Nothing
  case r1 of
    True -> do
      liftIO $ threadDelay 1000000
      r2 <- I.getServiceHealth ccManager (I.hostWithScheme _client) ccPort "testGetServiceHealth"
      case r2 of
        Just [x] -> return ()
        Just [] -> assertFailure "testGetServiceHealth: No Services Returned"
        Nothing -> assertFailure "testGetServiceHealth: Failed to parse result"
    False -> assertFailure "testGetServiceHealth: Service was not created"

testHealth :: TestTree
testHealth = testGroup "Health Check Tests" [testGetServiceHealth]

{- Session -}
testCreateSession :: TestTree
testCreateSession = testCase "testCreateSession" $ do
  _client@ConsulClient{..} <- client
  let req = SessionRequest Nothing (Just "testCreateSession") Nothing ["serfHealth"] (Just Release) (Just "30s")
  result <- I.createSession ccManager (I.hostWithScheme _client) ccPort req Nothing
  case result of
    Just _ -> return ()
    Nothing -> assertFailure "testCreateSession: No session was created"

testGetSessionInfo :: TestTree
testGetSessionInfo = testCase "testGetSessionInfo" $ do
  _client@ConsulClient{..} <- client
  let req = SessionRequest Nothing (Just "testGetSessionInfo") Nothing ["serfHealth"] (Just Release) (Just "30s")
  result <- I.createSession ccManager (I.hostWithScheme _client) ccPort req Nothing
  case result of
    Just x -> do
      x1 <- I.getSessionInfo ccManager (I.hostWithScheme _client) ccPort x Nothing
      case x1 of
        Just _ -> return ()
        Nothing -> assertFailure "testGetSessionInfo: Session Info was not returned"
    Nothing -> assertFailure "testGetSessionInfo: No session was created"

testRenewSession :: TestTree
testRenewSession = testCase "testRenewSession" $ do
  _client@ConsulClient{..} <- client
  let req = SessionRequest Nothing (Just "testRenewSession") Nothing ["serfHealth"] (Just Release) (Just "30s")
  result <- I.createSession ccManager (I.hostWithScheme _client) ccPort req Nothing
  case result of
    Just x -> do
      x1 <- I.renewSession ccManager (I.hostWithScheme _client) ccPort x Nothing
      case x1 of
        True -> return ()
        False -> assertFailure "testRenewSession: Session was not renewed"
    Nothing -> assertFailure "testRenewSession: No session was created"

testRenewNonexistentSession :: TestTree
testRenewNonexistentSession = testCase "testRenewNonexistentSession" $ do
  _client@ConsulClient{..} <- client
  sessId :: UUID <- randomIO
  let session = Session (toText sessId) Nothing
  x <- I.renewSession ccManager (I.hostWithScheme _client) ccPort session Nothing
  case x of
    True -> assertFailure "testRenewNonexistentSession: Non-existent session was renewed"
    False -> return ()

testDestroySession :: TestTree
testDestroySession = testCase "testDestroySession" $ do
  _client@ConsulClient{..} <- client
  let req = SessionRequest Nothing (Just "testDestroySession") Nothing ["serfHealth"] (Just Release) (Just "30s")
  result <- I.createSession ccManager (I.hostWithScheme _client) ccPort req Nothing
  case result of
    Just x -> do
      _ <- I.destroySession ccManager (I.hostWithScheme _client) ccPort x Nothing
      x1 <- I.getSessionInfo ccManager (I.hostWithScheme _client) ccPort x Nothing
      assertBool "testDestroySession: Session info was returned after destruction" $ (x1 == Nothing) || (x1 == Just [])
    Nothing -> assertFailure "testDestroySession: No session was created"

testInternalSession :: TestTree
testInternalSession = testGroup "Internal Session Tests" [testCreateSession, testGetSessionInfo, testRenewSession, testRenewNonexistentSession, testDestroySession]


testSessionMaintained :: TestTree
testSessionMaintained = testCase "testSessionMaintained" $ do
  client@ConsulClient{..} <- client
  let req = SessionRequest Nothing (Just "testSessionMaintained") Nothing ["serfHealth"] (Just Release) (Just "10s")
  result <- I.createSession ccManager (I.hostWithScheme client) ccPort req Nothing
  case result of
    Just session -> do
      threadDelay (12 * 1000000)
      y <- getSessionInfo client session Nothing
      assertEqual "testSessionMaintained: Session not found" True (isJust y)
    Nothing -> assertFailure "testSessionMaintained: No Session was created"


testWithSessionCancel :: TestTree
testWithSessionCancel = testCase "testWithSessionCancel" $ do
  client@ConsulClient{..} <- initializeConsulClient "localhost" 8500 Nothing
  let req = SessionRequest Nothing (Just "testWithSessionCancel") Nothing ["serfHealth"] (Just Release) (Just "10s")
  result <- I.createSession ccManager (I.hostWithScheme client) ccPort req Nothing
  case result of
    Just session -> do
      x1 <- withSession client Nothing 5 session (\ y -> action y client ) cancelAction
      assertEqual "testWithSessionCancel: Incorrect value" "Canceled" x1
      z <- getSessionInfo client session Nothing
      assertBool "testWithSessionCancel: Session was found" $ (z == Nothing) || (z == Just [])

    Nothing -> assertFailure "testWithSessionCancel: No session was created"
  where
    action :: MonadIO m => Session -> ConsulClient -> m Text
    action x client@ConsulClient{..} = do
      I.destroySession ccManager (I.hostWithScheme client) ccPort x Nothing
      liftIO $ threadDelay (30 * 1000000)
      return ("NotCanceled" :: Text)

    cancelAction :: MonadIO m => m Text
    cancelAction = return ("Canceled" :: Text)


testRunServiceTtl :: TestTree
testRunServiceTtl = testCase "testRunServiceTtl" $ do
  client@ConsulClient{..} <- initializeConsulClient "localhost" 8500 Nothing
  let register = RegisterService Nothing "testRunServiceTtl" [] (Just 8000) $ Just $ Ttl "10s"
  runService client register (action client) Nothing
  where
    action client = do
      threadDelay 15000000
      mHealth <- getServiceHealth client "testRunServiceTtl"
      case mHealth of
        Nothing -> assertFailure "testRunServiceTtl: No healthcheck was found"
        Just [x] -> do
          let checks = hChecks x
          mapM_ (testCheck) checks
    testCheck check = do
      assertBool "testRunServiceTtl: Check not passing" $ cStatus check == Passing


sessionWorkflowTests :: TestTree
sessionWorkflowTests = testGroup "Session Workflow Tests" [testWithSessionCancel,testSessionMaintained]

runServiceTests :: TestTree
runServiceTests = testGroup "Run Service Tests" [testRunServiceTtl]

agentTests :: TestTree
agentTests = testGroup "Agent Tests" [testGetSelf,testRegisterService]

allTests :: TestTree
allTests = testGroup "All Tests" [testInternalSession, internalKVTests, sessionWorkflowTests, agentTests,testHealth, clientKVTests, runServiceTests]

main :: IO ()
main = defaultMain allTests