module TAP
-- [todo] extract this to its own package?

import Options
import public Control.Monad.Either
import Data.String
import Data.List
import Data.List.Elem
import Data.SnocList
import Control.Monad.Reader
import Control.Monad.State
import Control.ANSI
import System

%default total


public export
Info : Type
Info = List (String, String)

private
data Result = Tried Bool Info | Skip String | Todo String

private
record TestBase where
  constructor MakeTest
  label : String
  run   : IO Result


private
toLines1 : (String, String) -> List String
toLines1 (k, v) =
  let vs = lines v in
  if length vs == 1
     then ["\{k}: \{v}"]
     else "\{k}: |" :: map (indent 2) vs

private
toLines : Info -> List String
toLines [] = []
toLines xs = "---" :: concatMap toLines1 xs <+> ["..."]


public export interface ToInfo e where toInfo : e -> Info

export ToInfo () where toInfo () = []

export Show a => ToInfo (List (String, a)) where toInfo = map (map show)


export
data Test
= One TestBase
| Group String (List Test)
| Note String


export
isRealTest : Test -> Bool
isRealTest (One   _)   = True
isRealTest (Group _ _) = True
isRealTest (Note  _)   = False


private
result : ToInfo a => Bool -> a -> IO Result
result ok = pure . Tried ok . toInfo

private
lazyToIO : Lazy a -> IO a
lazyToIO val = primIO $ \w => MkIORes (force val) w

export
testIO : (ToInfo e, ToInfo a) => String -> EitherT e IO a -> Test
testIO label act = One $ MakeTest label $ do
  case !(runEitherT act) of
       Right val => result True  val
       Left  err => result False err

export
test : (ToInfo e, ToInfo a) => String -> Lazy (Either e a) -> Test
test label val = testIO label $ MkEitherT $ lazyToIO val

export
todoWith : String -> String -> Test
todoWith label reason = One $ MakeTest label $ pure $ Todo reason

export
todo : String -> Test
todo label = todoWith label ""

private
makeSkip : String -> String -> Test
makeSkip label reason = One $ MakeTest label $ pure $ Skip reason

export
skipWith : Test -> String -> Test
skipWith (One   t)   reason = makeSkip t.label reason
skipWith (Group l _) reason = makeSkip l       reason
skipWith (Note  n)   _      = Note n

export
skip : Test -> Test
skip test = skipWith test ""

export
testThrows : (ToInfo e, Show a) =>
             String -> (e -> Bool) -> Lazy (Either e a) -> Test
testThrows label p act = One $ MakeTest label $ do
  case !(lazyToIO act) of
       Left  err => if p err then result True () else result False err
       Right val => result False [("success", val)]

infix 1 :-
export
(:-) : String -> List Test -> Test
(:-) = Group

export
bailOut : Test
bailOut = One $ MakeTest "bail out" $ do
  putStrLn "Bail out!"
  exitFailure

export
note : String -> Test
note = Note



export
header : List Test -> String
header tests =
  let count = length $ filter isRealTest tests in
  "1..\{show count}"


private
withPrefix : SnocList String -> TestBase -> Test
withPrefix pfx = One . {label $= (makePrefix pfx ++)}
  where makePrefix = concatMap $ \s => "\{s} ⟫ "

mutual
  export
  flattenWith : SnocList String -> List Test -> List Test
  flattenWith pfx tests =
    concatMap (\t => flatten1With pfx (assert_smaller tests t)) tests

  export
  flatten1With : SnocList String -> Test -> List Test
  flatten1With pfx (One   t)    = [withPrefix pfx t]
  flatten1With pfx (Group x ts) = flattenWith (pfx :< x) ts
  flatten1With pfx (Note  n)    = [Note n]

export
flatten : List Test -> List Test
flatten = flattenWith [<]

export
flatten1 : Test -> List Test
flatten1 = flatten1With [<]


private
record RunnerEnv where
  constructor RE
  indent : Nat
  color  : Bool


private
Runner : Type -> Type
Runner = ReaderT RunnerEnv IO

private
putIndentLines : List String -> Runner ()
putIndentLines xs = traverse_ (putStrLn . indent (!ask).indent) xs

private
isOk : Bool -> String
isOk b = if b then "ok" else "not ok"

private
toBool : Result -> Bool
toBool (Tried ok _) = ok
toBool _            = True


private
numbered : (a -> Bool) -> List a -> List (Nat, a)
numbered p = go 1 where
  go : Nat -> List a -> List (Nat, a)
  go _ []        = []
  go i (x :: xs) =
    if p x then (i, x) :: go (S i) xs
           else (0, x) :: go i     xs


private
col : Color -> String -> Runner String
col c str = pure $ if (!ask).color then show $ colored c str else str

private
putColor : Color -> String -> Runner ()
putColor c str = putIndentLines [!(col c str)]

private
okCol : Bool -> Color
okCol True  = Green
okCol False = Red

private
putOk' : Color -> Bool -> Nat -> String -> Runner ()
putOk' c ok index label =
  putIndentLines [!(col c "\{isOk ok} \{show index}") ++ " - \{label}"]

private
putOk : Bool -> Nat -> String -> Runner ()
putOk ok = putOk' (okCol ok) ok

private
putVersion : TAPVersion -> Runner ()
putVersion ver = putColor Cyan "TAP version \{show ver}"

private
run1' : (Nat, TestBase) -> Runner Bool
run1' (index, test) = do
  res <- liftIO test.run
  case res of
       Tried ok info => do
         putOk ok index test.label
         local {indent $= plus 2} $ putIndentLines $ toLines info
       Skip reason =>
         putOk' Yellow True index "\{test.label}  # skip \{reason}"
       Todo reason =>
         putOk' Yellow True index "\{test.label}  # todo \{reason}"
  pure $ toBool res

mutual
  private
  run' : (Nat, Test) -> Runner Bool
  run' (index, One test) = run1' (index, test)
  run' (index, Group label tests) = do
    putIndentLines [!(col Magenta "# Subtest: ") ++ label]
    res <- local {indent $= plus 4} $ runList tests
    putOk res index label
    pure res
  run' (_, Note note) = do
    putIndentLines [!(col Magenta "# ") ++ note]
    pure True

  private
  runList : List Test -> Runner Bool
  runList tests = do
    putColor Cyan $ header tests
    let tests' = numbered isRealTest tests
    all id <$> traverse (\t => run' (assert_smaller tests t)) tests'


mutual
  export
  filterMatch : Maybe String -> List Test -> List Test
  filterMatch Nothing    tests = tests
  filterMatch (Just pat) tests =
    mapMaybe (\t => filterMatch1 pat (assert_smaller tests t)) tests

  export
  filterMatch1 : String -> Test -> Maybe Test
  filterMatch1 pat test@(One base) =
    guard (pat `isInfixOf` base.label) $> test
  filterMatch1 pat all@(Group label tests) =
    if pat `isInfixOf` label then Just all else
    case filterMatch (Just pat) tests of
         []  => Nothing
         res => Just $ Group label res
  filterMatch1 pat note@(Note _) = Just note


export
main : Options -> List Test -> IO ExitCode
main opts tests = do
  let tests = filterMatch opts.pattern $
              case opts.version of V13 => flatten tests; V14 => tests
  let act = do putVersion opts.version; runList tests
  pure $ if !(runReaderT (RE 0 opts.color) act)
            then ExitSuccess
            else ExitFailure 70