module Language where

import Control.Monad.State (evalState)
import Data.Foldable (find)
import qualified Micro.Ast as A
import Micro.Compiler
import qualified Micro.Error as E
import Micro.Parser (parseFromString)
import Test.HUnit
import Text.Parsec.Pos (newPos)

assertCompileAst :: String -> [A.Expr] -> Assertion
assertCompileAst input expected = do
  r <- return $ parseFromString input
  case r of
    Left e -> assertFailure $ show e
    Right ast -> do
      res <- return $ compileToAst ast
      case res of
        Left e -> assertFailure $ show e
        Right ast -> assertEqual "" expected ast

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

expectError :: String -> E.ErrorType -> Assertion
expectError input etyp = do
  r <- return $ parseFromString input
  case r of
    Left e -> assertFailure $ show e
    Right ast -> do
      res <- return $ compile ast
      case res of
        Left e -> case (find (\(E.Error t _ _) -> t == etyp) e) of
          Just _ -> return ()
          Nothing -> assertFailure $ "expected " ++ show etyp ++ " didn't happen, got instead:\n" ++ unlines (map (\(E.Error t _ _) -> show t) e)
        Right _ -> assertFailure "expected error, didn't happen"

testCases =
  [ TestLabel "parse module" $
      TestCase $
        assertAst "module main" [A.Module "main" $ newPos "-" 1 1],
    TestLabel "parse binary number" $
      TestCase $
        assertAst
          "module main\n0b100000;"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Num 32 $ newPos "-" 2 1
          ],
    TestLabel
      "parse a function"
      $ TestCase $
        assertAst
          "module main\n\
          \def fn() { }"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Func "fn" [] Nothing [] False False $ newPos "-" 2 1
          ],
    TestLabel
      "parse a function with parameters"
      $ TestCase $
        assertAst
          "module main\n\
          \def fn(a: u8) { }"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Func "fn" [("a", A.Type "u8", True, newPos "-" 2 8)] Nothing [] False False $ newPos "-" 2 1
          ],
    TestLabel
      "parse a function with return value"
      $ TestCase $
        assertAst
          "module main\n\
          \def fn(): u8 {\n\
          \return 1; }"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Func "fn" [] (Just $ A.Type "u8") [A.Return (Just $ A.Num 1 $ newPos "-" 3 8) $ newPos "-" 3 1] False False $ newPos "-" 2 1
          ],
    TestLabel
      "parse a function call"
      $ TestCase $
        assertAst
          "module main\n\
          \def fn() { }\n\
          \fn();"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Func "fn" [] Nothing [] False False $ newPos "-" 2 1,
            A.Call (A.Variable "fn" $ newPos "-" 3 1) [] $ newPos "-" 3 1
          ],
    TestLabel
      "parse a function call with arguments"
      $ TestCase $
        assertAst
          "module main\n\
          \def fn(a: u8) { }\n\
          \fn(10);"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Func "fn" [("a", A.Type "u8", True, newPos "-" 2 8)] Nothing [] False False $ newPos "-" 2 1,
            A.Call (A.Variable "fn" $ newPos "-" 3 1) [A.Num 10 $ newPos "-" 3 4] $ newPos "-" 3 1
          ],
    TestLabel
      "parse empty return on a function"
      $ TestCase $
        assertAst
          "module main\n\
          \def fn() {\n\
          \return; }"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Func "fn" [] Nothing [A.Return Nothing $ newPos "-" 3 1] False False $ newPos "-" 2 1
          ],
    TestLabel
      "parse a recursive function"
      $ TestCase $
        assertAst
          "module main\n\
          \def fn() {\n\
          \fn(); }"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Func "fn" [] Nothing [A.Call (A.Variable "fn" $ newPos "-" 3 1) [] $ newPos "-" 3 1] False False $ newPos "-" 2 1
          ],
    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 "-" 1 1,
            A.Func "fn1" [] Nothing [] False False $ newPos "-" 2 1,
            A.Func
              "fn2"
              [("f", A.FuncType [] Nothing, True, newPos "-" 3 9)]
              Nothing
              [ A.Call (A.Variable "f" $ newPos "-" 4 1) [] $ newPos "-" 4 1
              ]
              False
              False
              $ newPos "-" 3 1,
            A.Call (A.Variable "fn2" $ newPos "-" 5 1) [A.Variable "fn1" $ newPos "-" 5 5] $ newPos "-" 5 1
          ],
    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 "-" 1 1,
            A.Func
              "fn"
              [("f", A.FuncType [] Nothing, True, newPos "-" 2 8)]
              Nothing
              [A.Call (A.Variable "f" $ newPos "-" 3 1) [] $ newPos "-" 3 1]
              False
              False
              $ newPos "-" 2 1,
            A.Call (A.Variable "fn" $ newPos "-" 4 1) [A.Func "lambda@4,4" [] Nothing [] True True $ newPos "-" 4 4] $ newPos "-" 4 1
          ],
    TestLabel "parse a call to lambda" $
      TestCase $
        assertAst
          "module main\n\
          \() { }();"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Call (A.Func "lambda@2,1" [] Nothing [] True True $ newPos "-" 2 1) [] $ newPos "-" 2 1
          ],
    TestLabel "parse a variable declaration" $
      TestCase $
        assertAst
          "module main\n\
          \var a: u8 = 10;"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Var "a" (A.Type "u8") (A.Num 10 $ newPos "-" 2 13) False $ newPos "-" 2 5
          ],
    TestLabel "parse a private variable declaration" $
      TestCase $
        assertAst
          "module main\n\
          \private var a: u8 = 10;"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Var "a" (A.Type "u8") (A.Num 10 $ newPos "-" 2 21) True $ newPos "-" 2 13
          ],
    TestLabel "parse a group variable declaration" $
      TestCase $
        assertAst
          "module main\n\
          \var (a: u8 = 10,\n\
          \b: bool = true);"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Var "a" (A.Type "u8") (A.Num 10 $ newPos "-" 2 14) False $ newPos "-" 2 6,
            A.Var "b" (A.Type "bool") (A.Bool' True $ newPos "-" 3 11) False $ newPos "-" 3 1
          ],
    TestLabel "parse a group of private variable declaration" $
      TestCase $
        assertAst
          "module main\n\
          \private var (a: u8 = 10,\n\
          \b: bool = true);"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Var "a" (A.Type "u8") (A.Num 10 $ newPos "-" 2 22) True $ newPos "-" 2 14,
            A.Var "b" (A.Type "bool") (A.Bool' True $ newPos "-" 3 11) True $ newPos "-" 3 1
          ],
    TestLabel "parse assignation" $
      TestCase $
        assertAst
          "module main\n\
          \var a: u8 = 0;\n\
          \a = 10;"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Var "a" (A.Type "u8") (A.Num 0 $ newPos "-" 2 13) False $ newPos "-" 2 5,
            A.BinOp A.Assign (newPos "-" 3 5) (A.Variable "a" $ newPos "-" 3 1) (A.Num 10 $ newPos "-" 3 5)
          ],
    TestLabel "fold constant addition" $
      TestCase $
        assertCompileAst
          "module main\n\
          \1 + 2 + 3 + (5 + 5);"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Num 16 $ newPos "-" 2 13
          ],
    TestLabel "fold constant addition in return" $
      TestCase $
        assertCompileAst
          "module main\n\
          \def fn(): u8 { return 1 + 2; }"
          [ A.Module "main" $ newPos "-" 1 1,
            A.Func "fn" [] (Just $ A.Type "u8") [A.Return (Just $ A.Num 3 $ newPos "-" 2 27) $ newPos "-" 2 16] False False $ newPos "-" 2 1
          ]
  ]

testErrors =
  [ TestLabel "invalid return value (empty return)" $
      TestCase $
        expectError
          "module main\n\
          \def fn(): u8 { return; }"
          E.TypeError,
    TestLabel "invalid return value" $
      TestCase $
        expectError
          "module main\n\
          \def fn(): u16 { return 1; }"
          E.TypeError,
    TestLabel "return without function" $
      TestCase $
        expectError
          "module main\n\
          \return;"
          E.UnexpectedReturn,
    TestLabel "symbol already defined" $
      TestCase $
        expectError
          "module main\n\
          \def fn() { }\n\
          \def fn() { }"
          E.AlreadyDefined,
    TestLabel "parameter already defined" $
      TestCase $
        expectError
          "module main\n\
          \def fn(a: u8, a: u8) { }\n"
          E.AlreadyDefined,
    TestLabel "call on non callable" $
      TestCase $
        expectError
          "module main\n\
          \def fn(a: u8): u8 { return a(); }\n"
          E.NonCallable,
    TestLabel "undefined variable" $
      TestCase $
        expectError
          "module main\n\
          \def fn(a: u8): u8 { return undef; }\n"
          E.Undefined,
    TestLabel "lambdas can use local variables only" $
      TestCase $
        expectError
          "module main\n\
          \def fn(a: u8): () -> u8 {\n\
          \return (): u8 { return a; };\n\
          \}\n"
          E.Undefined,
    TestLabel "undefined type in function parameters" $
      TestCase $
        expectError
          "module main\n\
          \def fn(a: undef): bool {\n\
          \return true;\n\
          \}\n"
          E.UndefinedType,
    TestLabel "undefined type in function return type" $
      TestCase $
        expectError
          "module main\n\
          \def fn(): undef {\n\
          \return true;\n\
          \}\n"
          E.UndefinedType,
    TestLabel "already defined variable" $
      TestCase $
        expectError
          "module main\n\
          \var a: u8 = 1; var a: u8 = 1;\n"
          E.AlreadyDefined,
    TestLabel "undefined type in variable declaration" $
      TestCase $
        expectError
          "module main\n\
          \var a: undef = 1;\n"
          E.UndefinedType,
    TestLabel "type mismatch in variable declaration" $
      TestCase $
        expectError
          "module main\n\
          \var a: bool = 0;\n"
          E.TypeError,
    TestLabel "type mismatch in variable declaration (lambda)" $
      TestCase $
        expectError
          "module main\n\
          \var a: bool = (a:u8): u8 { return a; };\n"
          E.TypeError,
    TestLabel "invalid assignation target" $
      TestCase $
        expectError
          "module main\n\
          \def fn() { return; }\n\
          \fn() = 10;"
          E.InvalidTarget,
    TestLabel "type mismatch in assignation" $
      TestCase $
        expectError
          "module main\n\
          \var a: u8 = 0;\n\
          \a = false;"
          E.TypeError,
    TestLabel "type mismatch in binary operator" $
      TestCase $
        expectError
          "module main\n\
          \1 + false;\n"
          E.TypeError,
    TestLabel "div by 0" $
      TestCase $
        expectError
          "module main\n\
          \1 / 0;\n"
          E.InvalidOperation
  ]

language = testCases ++ testErrors