module Language where

import qualified Ast as A
import Compiler
import Control.Monad.State (evalState)
import qualified Error as E
import Lexer (scan)
import Parser (parse)
import Test.HUnit
import Text.Parsec (runParser)
import Text.Parsec.Pos (newPos)

assertAst :: String -> [A.Expr] -> Assertion
assertAst input expected = do
  r <- return $ runParser (scan parse) () "test" input
  case r of
    Left e -> assertFailure $ show e
    Right ast -> assertEqual "" expected ast

expectError :: String -> IO (Maybe E.Error)
expectError input = do
  r <- return $ runParser (scan parse) () "test" input
  case r of
    Left e -> assertFailure $ show e
    Right ast -> do
      res <- return $ evalState (compileAll ast) startState
      case res of
        Left (e : _) -> return $ Just e
        Right _ -> return $ Nothing

testCase1 =
  TestLabel "parse module" $
    TestCase $
      assertAst "module main" [A.Module "main" $ newPos "test" 1 1]

testCase2 =
  TestLabel
    "parse a function"
    $ TestCase $
      assertAst
        "module main\n\
        \def fn() { }"
        [ A.Module "main" $ newPos "test" 1 1,
          A.Func "fn" [] Nothing [] False False $ newPos "test" 2 1
        ]

testCase3 =
  TestLabel
    "parse a function with parameters"
    $ TestCase $
      assertAst
        "module main\n\
        \def fn(a: u8) { }"
        [ A.Module "main" $ newPos "test" 1 1,
          A.Func "fn" [("a", A.Type "u8", newPos "test" 2 8)] Nothing [] False False $ newPos "test" 2 1
        ]

testCase4 =
  TestLabel
    "parse a function with return value"
    $ TestCase $
      assertAst
        "module main\n\
        \def fn(): u8 {\n\
        \return 1; }"
        [ A.Module "main" $ newPos "test" 1 1,
          A.Func "fn" [] (Just $ A.Type "u8") [A.Return (Just $ A.Num 1 $ newPos "test" 3 8) $ newPos "test" 3 1] False False $ newPos "test" 2 1
        ]

testCase5 =
  TestLabel
    "parse a function call"
    $ TestCase $
      assertAst
        "module main\n\
        \def fn() { }\n\
        \fn();"
        [ A.Module "main" $ newPos "test" 1 1,
          A.Func "fn" [] Nothing [] False False $ newPos "test" 2 1,
          A.Call (A.Var "fn" $ newPos "test" 3 1) [] $ newPos "test" 3 1
        ]

testCase6 =
  TestLabel
    "parse a function call with arguments"
    $ TestCase $
      assertAst
        "module main\n\
        \def fn(a: u8) { }\n\
        \fn(10);"
        [ A.Module "main" $ newPos "test" 1 1,
          A.Func "fn" [("a", A.Type "u8", newPos "test" 2 8)] Nothing [] False False $ newPos "test" 2 1,
          A.Call (A.Var "fn" $ newPos "test" 3 1) [A.Num 10 $ newPos "test" 3 4] $ newPos "test" 3 1
        ]

testCase7 =
  TestLabel
    "parse empty return on a function"
    $ TestCase $
      assertAst
        "module main\n\
        \def fn() {\n\
        \return; }"
        [ A.Module "main" $ newPos "test" 1 1,
          A.Func "fn" [] Nothing [A.Return Nothing $ newPos "test" 3 1] False False $ newPos "test" 2 1
        ]

testCase8 =
  TestLabel
    "parse a recursive function"
    $ TestCase $
      assertAst
        "module main\n\
        \def fn() {\n\
        \fn(); }"
        [ A.Module "main" $ newPos "test" 1 1,
          A.Func "fn" [] Nothing [A.Call (A.Var "fn" $ newPos "test" 3 1) [] $ newPos "test" 3 1] False False $ newPos "test" 2 1
        ]

testCase9 =
  TestLabel
    "parse a function with a function parameter"
    $ TestCase $
      assertAst
        "module main\n\
        \def fn1() { }\n\
        \def fn2(f: ()) {\n\
        \f(); }\n\
        \fn2(fn1);"
        [ A.Module "main" $ newPos "test" 1 1,
          A.Func "fn1" [] Nothing [] False False $ newPos "test" 2 1,
          A.Func
            "fn2"
            [("f", A.FuncType [] Nothing, newPos "test" 3 9)]
            Nothing
            [ A.Call (A.Var "f" $ newPos "test" 4 1) [] $ newPos "test" 4 1
            ]
            False
            False
            $ newPos "test" 3 1,
          A.Call (A.Var "fn2" $ newPos "test" 5 1) [A.Var "fn1" $ newPos "test" 5 5] $ newPos "test" 5 1
        ]

testCase10 =
  TestLabel
    "parse a function with a function parameter (lambda)"
    $ TestCase $
      assertAst
        "module main\n\
        \def fn(f: ()) {\n\
        \f(); }\n\
        \fn(() { });"
        [ A.Module "main" $ newPos "test" 1 1,
          A.Func
            "fn"
            [("f", A.FuncType [] Nothing, newPos "test" 2 8)]
            Nothing
            [A.Call (A.Var "f" $ newPos "test" 3 1) [] $ newPos "test" 3 1]
            False
            False
            $ newPos "test" 2 1,
          A.Call (A.Var "fn" $ newPos "test" 4 1) [A.Func "lambda@4,4" [] Nothing [] True True $ newPos "test" 4 4] $ newPos "test" 4 1
        ]

testCase11 =
  TestLabel "parse a call to lambda" $
    TestCase $
      assertAst
        "module main\n\
        \() { }();"
        [ A.Module "main" $ newPos "test" 1 1,
          A.Call (A.Func "lambda@2,1" [] Nothing [] True True $ newPos "test" 2 1) [] $ newPos "test" 2 1
        ]

testCase12 =
  TestLabel "invalid return value" $
    TestCase $ do
      e <-
        expectError
          "module main\n\
          \def fn(): u8 { return; }"
      case e of
        Nothing -> assertFailure "expected error, didn't happen"
        Just (E.TypeError _ _) -> return $ ()

testCase13 =
  TestLabel "return without function" $
    TestCase $ do
      e <-
        expectError
          "module main\n\
          \return;"
      case e of
        Nothing -> assertFailure "expected error, didn't happen"
        Just (E.UnexpectedReturn _ _) -> return $ ()

language =
  [ testCase1,
    testCase2,
    testCase3,
    testCase4,
    testCase5,
    testCase6,
    testCase7,
    testCase8,
    testCase9,
    testCase10,
    testCase11,
    testCase12,
    testCase13
  ]