{-# LANGUAGE OverloadedStrings #-}

module Felix.Test.Unit.Token (unitTests) where

import Base
import Felix.Report.Location
import Felix.Syntax.Adapt
import Felix.Syntax.Abstract (Associativity(..))
import Felix.Syntax.Interface
import Felix.Syntax.Pragma
import Felix.Syntax.Token

import Data.Text qualified as Text
import Test.Tasty
import Test.Tasty.HUnit
import Text.Megaparsec (errorBundlePretty)

unitTests :: TestTree
unitTests = testGroup "Lexer"
    [ testCase "nested math inside text returns to text" nestedMathInsideText
    , testCase "nested text and math returns to the enclosing frame" deeperAlternation
    , testCase "braces inside text do not close the text frame" textBraceNesting
    , testCase "cases is tokenized as an environment only inside math" casesOnlyInsideMath
    , testCase "imports retain source locations and POSIX spellings" locatedImports
    , testCase "commented environment starts are ignored"
        ignoresCommentedEnvironmentStart
    , testCase "empty inputs construct empty lexical syntax"
        constructsEmptyLexicalSyntax
    , testCase "extracts exact source fixity pragmas"
        extractsSourceFixityPragmas
    , testCase "rejects malformed reserved pragma lines"
        rejectsMalformedPragmas
    ]

nestedMathInsideText :: Assertion
nestedMathInsideText = do
    tokens <- tokensInProof "$\\text{if $x \\in A$ then}$"
    tokens `shouldBe`
        [ BeginEnv "proof"
        , BeginEnv "math"
        , BeginEnv "text"
        , Word "if"
        , BeginEnv "math"
        , Variable "x"
        , Command "in"
        , Variable "A"
        , EndEnv "math"
        , Word "then"
        , EndEnv "text"
        , EndEnv "math"
        , EndEnv "proof"
        ]

deeperAlternation :: Assertion
deeperAlternation = do
    tokens <- tokensInProof "$\\text{a $ \\text{b $c$ d} e$ f}$"
    tokens `shouldBe`
        [ BeginEnv "proof"
        , BeginEnv "math"
        , BeginEnv "text"
        , Word "a"
        , BeginEnv "math"
        , BeginEnv "text"
        , Word "b"
        , BeginEnv "math"
        , Variable "c"
        , EndEnv "math"
        , Word "d"
        , EndEnv "text"
        , Variable "e"
        , EndEnv "math"
        , Word "f"
        , EndEnv "text"
        , EndEnv "math"
        , EndEnv "proof"
        ]

textBraceNesting :: Assertion
textBraceNesting = do
    tokens <- tokensInProof "$\\text{a {b}}$"
    tokens `shouldBe`
        [ BeginEnv "proof"
        , BeginEnv "math"
        , BeginEnv "text"
        , Word "a"
        , InvisibleBraceL
        , Word "b"
        , InvisibleBraceR
        , EndEnv "text"
        , EndEnv "math"
        , EndEnv "proof"
        ]

casesOnlyInsideMath :: Assertion
casesOnlyInsideMath = do
    tokensInsideMath <- tokensInProof "$\\begin{cases}x\\end{cases}$"
    tokensInsideMath `shouldBe`
        [ BeginEnv "proof"
        , BeginEnv "math"
        , BeginEnv "cases"
        , Variable "x"
        , EndEnv "cases"
        , EndEnv "math"
        , EndEnv "proof"
        ]

    tokensOutsideMath <- tokensInProof "\\begin{cases}x\\end{cases}"
    assertBool
        "cases should not be tokenized as an environment outside math"
        (BeginEnv "cases" `notElem` tokensOutsideMath && EndEnv "cases" `notElem` tokensOutsideMath)

locatedImports :: Assertion
locatedImports = do
    let raw = Text.unlines
            [ "% heading"
            , "\\import{set/base.tex}"
            , "\\import{set\\special.tex}"
            , "\\begin{axiom}"
            , "  $x = x$."
            , "\\end{axiom}"
            ]
    case gatherImports (FileId maxBound) "import-unit" raw of
        Left err ->
            assertFailure (errorBundlePretty err)
        Right imports@[firstImport, secondImport] -> do
            assertEqual
                "import paths"
                ["set/base.tex", "set\\special.tex"]
                (unLocated <$> imports)
            assertEqual "first import line" 2 (locLine (startPos firstImport))
            assertEqual "second import line" 3 (locLine (startPos secondImport))
        Right imports ->
            assertFailure ("expected two imports, got " <> show imports)

ignoresCommentedEnvironmentStart :: Assertion
ignoresCommentedEnvironmentStart = do
    let raw = Text.unlines
            [ "% \\begin{signature}"
            , "ordinary text"
            , "\\begin{struct}"
            , "  an ordered set $X$ is a onesorted structure."
            , "\\end{struct}"
            ]
    (_, chunks) <-
        either
            (assertFailure . errorBundlePretty)
            pure
            (runLexer (FileId maxBound) "commented-environment" raw)
    case chunks of
        [Located{unLocated = BeginEnv "struct"} : _] -> pure ()
        _ -> assertFailure ("expected the real structure environment, got " <> show chunks)

constructsEmptyLexicalSyntax :: Assertion
constructsEmptyLexicalSyntax =
    forM_
        [ ("empty", "")
        , ("comment-only", "% heading\n% body")
        ]
        \(description, raw) -> do
            let input = Text.pack raw
            imports <-
                either
                    (assertFailure . errorBundlePretty)
                    pure
                    (gatherImports
                        (FileId maxBound)
                        description
                        input)
            assertEqual (description <> " imports") [] imports
            (lexedImports, chunks) <-
                either
                    (assertFailure . errorBundlePretty)
                    pure
                    (runLexer
                        (FileId maxBound)
                        description
                        input)
            assertEqual (description <> " lexer imports") [] lexedImports
            assertEqual (description <> " chunks") [] chunks
            scanned <-
                either
                    (assertFailure . show)
                    pure
                    (concat <$> traverse scanChunk chunks)
            assertEqual (description <> " scanned declarations") [] scanned
            delta <-
                either
                    (assertFailure . show)
                    pure
                    (canonicalSyntaxDelta [])
            assertBool
                (description <> " syntax declarations")
                (null (canonicalSyntaxDeltaEntries delta))

extractsSourceFixityPragmas :: Assertion
extractsSourceFixityPragmas = do
    let input =
            "  %! infixl 0\n"
                <> "\t%! infixr 07\r\n"
                <> "%! infix 3"
        ordinaryComments =
            Text.unlines
                [ "% ! infixl 1"
                , "text %! infixr 2"
                , "% ordinary"
                ]
    pragmas <-
        either
            (assertFailure . Text.unpack . renderSyntaxPragmaError)
            pure
            (extractSyntaxPragmas
                (FileId maxBound)
                "pragma-unit"
                input)
    assertEqual
        "normalized pragmas"
        [ (LeftAssoc, 0, 1, 3)
        , (RightAssoc, 7, 2, 2)
        , (NonAssoc, 3, 3, 1)
        ]
        [ ( syntaxPragmaAssociativity pragma
          , sourceMixfixLevelValue (syntaxPragmaLevel pragma)
          , locLine (syntaxPragmaLocation pragma)
          , locColumn (syntaxPragmaLocation pragma)
          )
        | pragma <- pragmas
        ]
    assertEqual
        "ordinary comments"
        (Right [])
        (extractSyntaxPragmas
            (FileId maxBound)
            "pragma-unit"
            ordinaryComments)

rejectsMalformedPragmas :: Assertion
rejectsMalformedPragmas =
    forM_
        [ ("%!infixl 1\n", SyntaxPragmaMissingSpaceAfterPrefix)
        , ("%!\n", SyntaxPragmaMissingKeyword)
        , ("%! Infixl 1\n", SyntaxPragmaUnknownKeyword "Infixl")
        , ("%! infixl\n", SyntaxPragmaMissingLevel)
        , ("%! infixl -1\n", SyntaxPragmaInvalidLevel)
        , ("%! infixl ١\n", SyntaxPragmaInvalidLevel)
        , ("%! infixl 8\n", SyntaxPragmaLevelOutOfRange)
        , ("%! infixl 1 extra\n", SyntaxPragmaTrailingContent)
        , ("%! infixl 1\r", SyntaxPragmaLoneCarriageReturn)
        ]
        \(input, expectedProblem) ->
            case extractSyntaxPragmas
                    (FileId maxBound)
                    "pragma-unit"
                    input of
                Left (InvalidSyntaxPragma location actualProblem) -> do
                    assertEqual
                        ("problem for " <> show input)
                        expectedProblem
                        actualProblem
                    assertEqual "error line" 1 (locLine location)
                    assertEqual "error column" 1 (locColumn location)
                Left err ->
                    assertFailure
                        ("unexpected pragma error: "
                            <> Text.unpack (renderSyntaxPragmaError err))
                Right pragmas ->
                    assertFailure
                        ("expected malformed pragma rejection, got "
                            <> show pragmas)

tokensInProof :: Text -> IO [Token]
tokensInProof raw =
    case runLexer (FileId maxBound) "lexer-unit" wrapped of
        Left err ->
            assertFailure (errorBundlePretty err)
        Right (_imports, chunks) ->
            pure (concatMap (map unLocated) chunks)
    where
        wrapped = Text.unlines
            [ "\\begin{proof}"
            , raw
            , "\\end{proof}"
            ]

shouldBe :: (Eq a, Show a, HasCallStack) => a -> a -> Assertion
shouldBe = flip (assertEqual "")