--
-- Licensed to the Apache Software Foundation (ASF) under one
-- or more contributor license agreements. See the NOTICE file
-- distributed with this work for additional information
-- regarding copyright ownership. The ASF licenses this file
-- to you under the Apache License, Version 2.0 (the
-- "License"); you may not use this file except in compliance
-- with the License. You may obtain a copy of the License at
--
--   http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing,
-- software distributed under the License is distributed on an
-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-- KIND, either express or implied. See the License for the
-- specific language governing permissions and limitations
-- under the License.
--

{-# LANGUAGE OverloadedStrings, RecordWildCards, ScopedTypeVariables #-}
module Main where

import Control.Exception
import Control.Monad
import Data.Functor
import Data.List.Split
import Data.String
import Network
import Network.URI
import System.Environment
import System.Exit
import qualified Data.ByteString.Lazy as LBS
import qualified Data.HashMap.Strict as Map
import qualified Data.HashSet as Set
import qualified Data.Vector as Vector
import qualified System.IO as IO

import ThriftTest_Iface
import ThriftTest_Types
import qualified ThriftTest_Client as Client

import Thrift.Transport
import Thrift.Transport.Framed
import Thrift.Transport.Handle
import Thrift.Transport.HttpClient
import Thrift.Protocol
import Thrift.Protocol.Binary
import Thrift.Protocol.Compact
import Thrift.Protocol.Header
import Thrift.Protocol.JSON

data Options = Options
  { host         :: String
  , port         :: Int
  , domainSocket :: String
  , transport    :: String
  , protocol     :: ProtocolType
  -- TODO: Haskell lib does not have SSL support
  , ssl          :: Bool
  , testLoops    :: Int
  }
  deriving (Show, Eq)

data TransportType = Buffered IO.Handle
                   | Framed (FramedTransport IO.Handle)
                   | Http HttpClient
                   | NoTransport String

getTransport :: String -> String -> Int -> (IO TransportType)
getTransport "buffered" host port = do
  h <- hOpen (host, PortNumber $ fromIntegral port)
  IO.hSetBuffering h $ IO.BlockBuffering Nothing
  return $ Buffered h
getTransport "framed" host port = do
  h <- hOpen (host, PortNumber $ fromIntegral port)
  t <- openFramedTransport h
  return $ Framed t
getTransport "http" host port = let uriStr = "http://" ++ host ++ ":" ++ show port in
                                case parseURI uriStr of
                                  Nothing -> do return (NoTransport $ "Failed to parse URI: " ++ uriStr)
                                  Just(uri) -> do
                                    t <- openHttpClient uri
                                    return $ Http t
getTransport t host port = do return (NoTransport $ "Unsupported transport: " ++ t)

data ProtocolType = Binary
                  | Compact
                  | JSON
                  | Header
                  deriving (Show, Eq)

getProtocol :: String -> ProtocolType
getProtocol "binary"  = Binary
getProtocol "compact" = Compact
getProtocol "json"    = JSON
getProtocol "header" = Header
getProtocol p = error $ "Unsupported Protocol: " ++ p

defaultOptions :: Options
defaultOptions = Options
  { port         = 9090
  , domainSocket = ""
  , host         = "localhost"
  , transport    = "buffered"
  , protocol     = Binary
  , ssl          = False
  , testLoops    = 1
  }

runClient :: Protocol p => p -> IO ()
runClient p = do
  let prot = (p,p)
  putStrLn "Starting Tests"

  -- VOID Test
  putStrLn "testVoid"
  Client.testVoid prot

  -- String Test
  putStrLn "testString"
  s <- Client.testString prot "Test"
  when (s /= "Test") exitFailure

  -- Bool Test
  putStrLn "testBool"
  bool <- Client.testBool prot True
  when (not bool) exitFailure
  putStrLn "testBool"
  bool <- Client.testBool prot False
  when (bool) exitFailure

  -- Byte Test
  putStrLn "testByte"
  byte <- Client.testByte prot 1
  when (byte /= 1) exitFailure

  -- I32 Test
  putStrLn "testI32"
  i32 <- Client.testI32 prot (-1)
  when (i32 /= -1) exitFailure

  -- I64 Test
  putStrLn "testI64"
  i64 <- Client.testI64 prot (-34359738368)
  when (i64 /= -34359738368) exitFailure

  -- Double Test
  putStrLn "testDouble"
  dub <- Client.testDouble prot (-5.2098523)
  when (abs (dub + 5.2098523) > 0.001) exitFailure

  -- Binary Test
  putStrLn "testBinary"
  bin <- Client.testBinary prot (LBS.pack . reverse $ [-128..127])
  when ((reverse [-128..127]) /= LBS.unpack bin) exitFailure
  
  -- Struct Test
  let structIn = Xtruct{ xtruct_string_thing = "Zero"
                       , xtruct_byte_thing   = 1
                       , xtruct_i32_thing    = -3
                       , xtruct_i64_thing    = -5
                       }
  putStrLn "testStruct"
  structOut <- Client.testStruct prot structIn
  when (structIn /= structOut) exitFailure

  -- Nested Struct Test
  let nestIn = Xtruct2{ xtruct2_byte_thing   = 1
                      , xtruct2_struct_thing = structIn
                      , xtruct2_i32_thing    = 5
                      }
  putStrLn "testNest"
  nestOut <- Client.testNest prot nestIn
  when (nestIn /= nestOut) exitFailure

  -- Map Test
  let mapIn = Map.fromList $ map (\i -> (i, i-10)) [1..5]
  putStrLn "testMap"
  mapOut <- Client.testMap prot mapIn
  when (mapIn /= mapOut) exitFailure

  -- Set Test
  let setIn = Set.fromList [-2..3]
  putStrLn "testSet"
  setOut <- Client.testSet prot setIn
  when (setIn /= setOut) exitFailure

  -- List Test
  let listIn = Vector.fromList [-2..3]
  putStrLn "testList"
  listOut <- Client.testList prot listIn
  when (listIn /= listOut) exitFailure

  -- Enum Test
  putStrLn "testEnum"
  numz1 <- Client.testEnum prot Numberz_ONE
  when (numz1 /= Numberz_ONE) exitFailure

  putStrLn "testEnum"
  numz2 <- Client.testEnum prot Numberz_TWO
  when (numz2 /= Numberz_TWO) exitFailure

  putStrLn "testEnum"
  numz5 <- Client.testEnum prot Numberz_FIVE
  when (numz5 /= Numberz_FIVE) exitFailure

  -- Typedef Test
  putStrLn "testTypedef"
  uid <- Client.testTypedef prot 309858235082523
  when (uid /= 309858235082523) exitFailure

  -- Nested Map Test
  putStrLn "testMapMap"
  _ <- Client.testMapMap prot 1

  -- Exception Test
  putStrLn "testException"
  exn1 <- try $ Client.testException prot "Xception"
  case exn1 of
    Left (Xception _ _) -> return ()
    _ -> putStrLn (show exn1) >> exitFailure

  putStrLn "testException"
  exn2 <- try $ Client.testException prot "TException"
  case exn2 of
    Left (_ :: SomeException) -> return ()
    Right _ -> exitFailure

  putStrLn "testException"
  exn3 <- try $ Client.testException prot "success"
  case exn3 of
    Left (_ :: SomeException) -> exitFailure
    Right _ -> return ()

  -- Multi Exception Test
  putStrLn "testMultiException"
  multi1 <- try $ Client.testMultiException prot "Xception" "test 1"
  case multi1 of
    Left (Xception _ _) -> return ()
    _ -> exitFailure

  putStrLn "testMultiException"
  multi2 <- try $ Client.testMultiException prot "Xception2" "test 2"
  case multi2 of
    Left (Xception2 _ _) -> return ()
    _ -> exitFailure

  putStrLn "testMultiException"
  multi3 <- try $ Client.testMultiException prot "success" "test 3"
  case multi3 of
    Left (_ :: SomeException) -> exitFailure
    Right _ -> return ()


main :: IO ()
main = do
  options <- flip parseFlags defaultOptions <$> getArgs
  case options of
    Nothing -> showHelp
    Just Options{..} -> do
      trans <- Main.getTransport transport host port
      case trans of
        Buffered t -> runTest testLoops protocol t
        Framed t   -> runTest testLoops protocol t
        Http t     -> runTest testLoops protocol t
        NoTransport err -> putStrLn err
  where
    makeClient p t = case p of
                       Binary  -> runClient $ BinaryProtocol t
                       Compact -> runClient $ CompactProtocol t
                       JSON    -> runClient $ JSONProtocol t
                       Header  -> createHeaderProtocol t t >>= runClient
    runTest loops p t = do
      let client = makeClient p t
      replicateM_ loops client
      putStrLn "COMPLETED SUCCESSFULLY"

parseFlags :: [String] -> Options -> Maybe Options
parseFlags (flag : flags) opts = do
  let pieces = splitOn "=" flag
  case pieces of
    "--port" : arg : _ -> parseFlags flags opts{ port = read arg }
    "--domain-socket" : arg : _ -> parseFlags flags opts{ domainSocket = read arg }
    "--host" : arg : _ -> parseFlags flags opts{ host = arg }
    "--transport" : arg : _ -> parseFlags flags opts{ transport = arg }
    "--protocol" : arg : _ -> parseFlags flags opts{ protocol = getProtocol arg }
    "-n" : arg : _ -> parseFlags flags opts{ testLoops = read arg }
    "--h" : _ -> Nothing
    "--help" : _ -> Nothing
    "--ssl" : _ -> parseFlags flags opts{ ssl = True }
    "--processor-events" : _ -> parseFlags flags opts
    _ -> Nothing
parseFlags [] opts = Just opts

showHelp :: IO ()
showHelp = putStrLn
  "Allowed options:\n\
  \  -h [ --help ]               produce help message\n\
  \  --host arg (=localhost)     Host to connect\n\
  \  --port arg (=9090)          Port number to connect\n\
  \  --domain-socket arg         Domain Socket (e.g. /tmp/ThriftTest.thrift),\n\
  \                              instead of host and port\n\
  \  --transport arg (=buffered) Transport: buffered, framed, http\n\
  \  --protocol arg (=binary)    Protocol: binary, compact, json\n\
  \  --ssl                       Encrypted Transport using SSL\n\
  \  -n [ --testloops ] arg (=1) Number of Tests"