module Main where

import Eval
import Lexer
import Parser
import Research
import Control.Exception (evaluate, try, SomeException)
import qualified Data.Map as Map
import Test.Tasty
import Test.Tasty.HUnit
import Test.Tasty.QuickCheck
import Text.Megaparsec (runParser)

main :: IO ()
main = defaultMain tests

tests :: TestTree
tests = testGroup "Sapling Tests"
    [ lexerTests
    , parserTests
    , integrationTests
    , evaluationTests
    , propertyTests
    ]

lexerTests :: TestTree
lexerTests = testGroup "Lexer Tests"
    [ testCase "Lex simple identifiers" $ do
        let input = "x a b = a"
            expect = Right [LIdentifier "x", LIdentifier "a", LIdentifier "b", LAssign, LIdentifier "a"]
        runParser saplingLexer "" input @?= expect
    , testCase "Lex Tree Calculus terms" $ do
        let input = "t t t"
            expect = Right [LKeywordT, LKeywordT, LKeywordT]
        runParser saplingLexer "" input @?= expect
    , testCase "Lex escaped characters in strings" $ do
        let input = "\"hello\\nworld\""
            expect = Right [LStringLiteral "hello\\nworld"]
        runParser saplingLexer "" input @?= expect
    , testCase "Lex mixed literals" $ do
        let input = "t \"string\" 42"
            expect = Right [LKeywordT, LStringLiteral "string", LIntegerLiteral 42]
        runParser saplingLexer "" input @?= expect
    , testCase "Lex invalid token" $ do
        let input = "$invalid"
        case runParser saplingLexer "" input of
            Left _ -> return ()
            Right _ -> assertFailure "Expected lexer to fail on invalid token"
    , testCase "Drop trailing whitespace in definitions" $ do
        let input = "x = 5 "
            expect = [LIdentifier "x",LAssign,LIntegerLiteral 5]
        case (runParser saplingLexer "" input) of
            Left _ -> assertFailure "Failed to lex input"
            Right i -> i @?= expect
    , testCase "Error when using invalid characters in identifiers" $ do
        case (runParser saplingLexer "" "__result = 5") of
            Left _ -> return ()
            Right _ -> assertFailure "Expected failure when trying to assign the value of __result"
    ]

parserTests :: TestTree
parserTests = testGroup "Parser Tests"
    [ testCase "Error when parsing incomplete definitions" $ do
        let input = lexSapling "x = "
        case (runParser parseExpression "" input) of
            Left _ -> return ()
            Right _ -> assertFailure "Expected failure on invalid input"
    , testCase "Error when assigning a value to T" $ do
        let input = lexSapling "t = x"
        case (runParser parseExpression "" input) of
            Left _ -> return ()
            Right _ -> assertFailure "Expected failure when trying to assign the value of T"
    , testCase "Error when parsing bodyless definitions with arguments" $ do
        let input = lexSapling "x a b = "
        case (runParser parseExpression "" input) of
            Left _ -> return ()
            Right _ -> assertFailure "Expected failure on invalid input"
    , testCase "Parse function definitions" $ do
        let input = "x a b c = a"
        let expect = SFunc "x" ["a","b","c"] (SVar "a")
        parseSingle input @?= expect
    , testCase "Parse nested Tree Calculus terms" $ do
        let input = "t (t t) t"
        let expect = SApp (SApp TLeaf (SApp TLeaf TLeaf)) TLeaf
        parseSingle input @?= expect
    , testCase "Parse sequential Tree Calculus terms" $ do
        let input = "t t t"
        let expect = SApp (SApp TLeaf TLeaf) TLeaf
        parseSingle input @?= expect
    , testCase "Parse mixed list literals" $ do
        let input = "[t (\"hello\") t]"
        let expect = SList [TLeaf, SStr "hello", TLeaf]
        parseSingle input @?= expect
    , testCase "Parse function with applications" $ do
        let input = "f x = t x"
        let expect = SFunc "f" ["x"] (SApp TLeaf (SVar "x"))
        parseSingle input @?= expect
     , testCase "Parse nested lists" $ do
        let input = "[t [(t t)]]"
        let expect = SList [TLeaf,SList [SApp TLeaf TLeaf]]
        parseSingle input @?= expect
    , testCase "Parse complex parentheses" $ do
        let input = "t (t t (t t))"
        let expect = SApp TLeaf (SApp (SApp TLeaf TLeaf) (SApp TLeaf TLeaf))
        parseSingle input @?= expect
    , testCase "Parse empty list" $ do
        let input = "[]"
        let expect = SList []
        parseSingle input @?= expect
    , testCase "Parse multiple nested lists" $ do
        let input = "[[t t] [t (t t)]]"
        let expect = SList [SList [TLeaf,TLeaf],SList [TLeaf,SApp TLeaf TLeaf]]
        parseSingle input @?= expect
    , testCase "Parse whitespace variance" $ do
        let input1 = "[t t]"
        let input2 = "[ t t ]"
        let expect = SList [TLeaf, TLeaf]
        parseSingle input1 @?= expect
        parseSingle input2 @?= expect
    , testCase "Parse string in list" $ do
        let input = "[(\"hello\")]"
        let expect = SList [SStr "hello"]
        parseSingle input @?= expect
    , testCase "Parse parentheses inside list" $ do
        let input = "[t (t t)]"
        let expect = SList [TLeaf,SApp TLeaf TLeaf]
        parseSingle input @?= expect
    , testCase "Parse nested parentheses in function body" $ do
        let input = "f = t (t (t t))"
        let expect = SFunc "f" [] (SApp TLeaf (SApp TLeaf (SApp TLeaf TLeaf)))
        parseSingle input @?= expect
    , testCase "Parse lambda abstractions" $ do
        let input = "(\\a : a)"
        let expect = (SLambda ["a"] (SVar "a"))
        parseSingle input @?= expect
    , testCase "Parse multiple arguments to lambda abstractions" $ do
        let input = "x = (\\a b : a)"
        let expect = SFunc "x" [] (SLambda ["a"] (SLambda ["b"] (SVar "a")))
        parseSingle input @?= expect
     , testCase "Grouping T terms with parentheses in function application" $ do
        let input  = "x = (\\a : a)\n" <> "x (t)"
            expect = [SFunc "x" [] (SLambda ["a"] (SVar "a")),SApp (SVar "x") TLeaf]
        parseSapling input @?= expect
    ]

integrationTests :: TestTree
integrationTests = testGroup "Integration Tests"
    [ testCase "Combine lexer and parser" $ do
        let input = "x = t t t"
        let expect = SFunc "x" [] (SApp (SApp TLeaf TLeaf) TLeaf)
        parseSingle input @?= expect
    , testCase "Complex Tree Calculus expression" $ do
        let input = "t (t t t) t"
        let expect = SApp (SApp TLeaf (SApp (SApp TLeaf TLeaf) TLeaf)) TLeaf
        parseSingle input @?= expect
    ]

evaluationTests :: TestTree
evaluationTests = testGroup "Evaluation Tests"
    [ testCase "Evaluate single Leaf" $ do
        let input = "t"
        let ast = parseSingle input
        (result $ evalSingle Map.empty ast) @?= Leaf
    , testCase "Evaluate single Stem" $ do
        let input = "t t"
        let ast = parseSingle input
        (result $ evalSingle Map.empty ast) @?= Stem Leaf
    , testCase "Evaluate single Fork" $ do
        let input = "t t t"
        let ast = parseSingle input
        (result $ evalSingle Map.empty ast) @?= Fork Leaf Leaf
    , testCase "Evaluate nested Fork and Stem" $ do
        let input = "t (t t) t"
        let ast = parseSingle input
        (result $ evalSingle Map.empty ast) @?= Fork (Stem Leaf) Leaf
    , testCase "Evaluate `not` function" $ do
        let input = "t (t (t t) (t t t)) t"
        let ast = parseSingle input
        (result $ evalSingle Map.empty ast) @?=
            Fork (Fork (Stem Leaf) (Fork Leaf Leaf)) Leaf
    , testCase "Environment updates with definitions" $ do
        let input = "x = t\ny = x"
        let env = evalSapling Map.empty (parseSapling input)
        Map.lookup "x" env @?= Just Leaf
        Map.lookup "y" env @?= Just Leaf
    , testCase "Variable substitution" $ do
        let input = "x = t t\ny = t x\ny"
        let env = evalSapling Map.empty (parseSapling input)
        (result env) @?= Stem (Stem Leaf)
     , testCase "Multiline input evaluation" $ do
        let input = "x = t\ny = t t\nx"
        let env = evalSapling Map.empty (parseSapling input)
        (result env) @?= Leaf
     , testCase "Evaluate string literal" $ do
        let input = "\"hello\""
        let ast = parseSingle input
        (result $ evalSingle Map.empty ast) @?= toString "hello"
    , testCase "Evaluate list literal" $ do
        let input = "[t (t t)]"
        let ast = parseSingle input
        (result $ evalSingle Map.empty ast) @?= toList [Leaf, Stem Leaf]
    , testCase "Evaluate empty list" $ do
        let input = "[]"
        let ast = parseSingle input
        (result $ evalSingle Map.empty ast) @?= toList []
    , testCase "Evaluate variable dependency chain" $ do
        let input = "x = t (t t)\n \
                    \ y = x\n \
                    \ z = y\n \
                    \ variablewithamuchlongername = z\n \
                    \ variablewithamuchlongername"
        let env = evalSapling Map.empty (parseSapling input)
        (result env) @?= (Stem (Stem Leaf))
    , testCase "Evaluate variable shadowing" $ do
        let input = "x = t t\nx = t\nx"
        let env = evalSapling Map.empty (parseSapling input)
        (result env) @?= Leaf
    , testCase "Lambda identity" $ do
        let input = "(\\a : a)"
            env = evalSapling Map.empty (parseSapling input)
        result env @?= Fork (Stem (Stem Leaf)) (Stem Leaf)
    , testCase "Apply identity to Boolean Not" $ do
        let not = "(t (t (t t) (t t t)) t)"
            input = "x = (\\a : a)\nx " ++ not
            env = evalSapling Map.empty (parseSapling input)
        result env @?= Fork (Fork (Stem Leaf) (Fork Leaf Leaf)) Leaf
    , testCase "Constant function matches" $ do
        let input = "k = (\\a b : a)\nk (t t) t"
            env = evalSapling Map.empty (parseSapling input)
        result env @?= Stem Leaf
    ]

propertyTests :: TestTree
propertyTests = testGroup "Property Tests"
    [ testProperty "Lexing and parsing round-trip" $ \input ->
        case runParser saplingLexer "" input of
            Left _ -> property True
            Right tokens -> case runParser parseExpression "" tokens of
                Left _ -> property True
                Right ast -> parseSingle input === ast
    ]