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"
      let 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"
      let expect = Right [LKeywordT, LKeywordT, LKeywordT]
      runParser saplingLexer "" input @?= expect

  , testCase "Lex escaped characters in strings" $ do
      let input = "\"hello\\nworld\""
      let expect = Right [LStringLiteral "hello\\nworld"]
      runParser saplingLexer "" input @?= expect

  , testCase "Lex mixed literals" $ do
      let input = "t \"string\" 42"
      let 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 "
      let 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 = a"
      let expect = SFunc "x" ["a", "b"] (SApp (SVar "a") [])
      parseSingle input @?= expect

  , testCase "Parse nested Tree Calculus terms" $ do
      let input = "t (t t) t"
      let expect = TFork (TStem TLeaf) TLeaf
      parseSingle input @?= expect

  , testCase "Parse sequential Tree Calculus terms" $ do
      let input = "t t t"
      let expect = TFork 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 [TStem TLeaf]]
      parseSingle input @?= expect

  , testCase "Parse complex parentheses" $ do
      let input = "t (t t (t t))"
      let expect = TStem (TFork TLeaf (TStem 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, TStem 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, TStem TLeaf]
      parseSingle input @?= expect

  , testCase "Parse nested parentheses in function body" $ do
      let input = "f = t (t (t t))"
      let expect = SFunc "f" [] (TStem (TStem (TStem TLeaf)))
      parseSingle input @?= expect
  ]

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

  , testCase "Complex Tree Calculus expression" $ do
      let input = "t (t t t) t"
      let expect = TFork (TFork 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\n \
                 \ y = t x\n \
                 \ z = t y\n \
                 \ variablewithamuchlongername = z\n \
                 \ variablewithamuchlongername"
      let env = evalSapling Map.empty (parseSapling input)
      (result env) @?= (Stem (Stem Leaf))

  , testCase "Evaluate redefinition of variables" $ do
      let input = "x = t t\nx = t\nx"
      let env = evalSapling Map.empty (parseSapling input)
      (result env) @?= 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
  ]