{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}

module Minilang.Lambda.ParserSpec where

import Data.Either (isLeft, isRight)
import Data.Function ((&))
import Data.Text (Text, pack)
import qualified Data.Text as Text
import Minilang.Lambda.Eval (Term (Lam))
import qualified Minilang.Lambda.Eval as Eval
import Minilang.Lambda.Parser (AST (..), desugar, initialChars, parse, parseProgram, pretty, restChars)
import Minilang.Lambda.Eval (normalise)
import Test.Hspec (Spec, describe, it, parallel, shouldBe, shouldSatisfy)
import Test.Hspec.QuickCheck (prop)
import Test.QuickCheck (
  Arbitrary (..),
  Gen,
  NonEmptyList (..),
  Positive (..),
  counterexample,
  elements,
  forAll,
  forAllShrink,
  listOf,
  oneof,
  resize,
  sized,
  suchThat,
  vectorOf,
  (===),
  (==>),
 )

-- ((f x) y) --> (f x y)
-- (lam (x) (lam (y) y)) --> (lam (x y) y)

spec :: Spec
spec = parallel $ do
  prop "parses an identifier as a variable" $ \(Identifier ident) ->
    parse ident `shouldBe` Right (Sym ident)

  prop "parses a lambda-expression as an abstraction" $ \(Identifier ident) (Identifier body) ->
    parse ("(lam (" <> ident <> ") " <> body <> ")") `shouldBe` Right (Abs [ident] (Sym body))

  prop "parses a lambda-expression with multiple bindings as an abstraction" $ \(NonEmpty idents) (Identifier body) ->
    let vars = unIdent <$> idents
        abs = Abs vars (Sym body)
     in parse ("(lam (" <> Text.unwords vars <> ") " <> body <> ")") `shouldBe` Right abs

  prop "parses an application" $ \(Identifier ident1) (Identifier ident2) ->
    parse ("(" <> ident1 <> " " <> ident2 <> ")") `shouldBe` Right (App (Sym ident1) (Sym ident2) [])

  prop "parses multiple applications" $ \(NonEmpty idents) ->
    (length idents >= 2) ==>
      let vars = unIdent <$> idents
          (a : b : rest) = vars
          app = App (Sym a) (Sym b) (Sym <$> rest)
       in parse ("(" <> Text.unwords vars <> ")") `shouldBe` Right app

  prop "parses is inverse to pretty" $
    forAllShrink (oneof [genAst 5, genDef 5]) shrink $ \ast ->
      parse (pretty ast) === Right ast

  prop "desugars abstractions to λ-terms" $ forAll (genAbstraction 5) $ \ast ->
    let term = desugar ast
        ar = arity ast
        nestedAbs (Lam _ body) = 1 + nestedAbs body
        nestedAbs _ = 0
     in nestedAbs term === ar
          & counterexample ("Desugared term: " <> show term)

  prop "desugars applications to app-terms" $ forAllShrink (genApplication 5) shrink $ \ast ->
    let term = desugar ast
        ar = spine ast
        nestedApp (Eval.App l _) = 1 + nestedApp l
        nestedApp _ = 0
     in nestedApp term === ar
          & counterexample ("Desugared term: " <> show term)

  prop "parses multiple let-binding as a function application" $
    forAll arbitrary $ \(NonEmpty idents) ->
      forAll (vectorOf (length idents) $ genAst 3) $ \(binding : bindings) ->
        forAll (genAst 2) $ \body ->
          let fun = Abs (unIdent <$> idents) body
              app = App fun binding bindings
              bound = zip idents (binding : bindings)
           in parse
                ( "(let ("
                    <> Text.concat
                      ( map
                          ( \(Identifier v, t) ->
                              "(" <> v <> " " <> pretty t <> ")"
                          )
                          bound
                      )
                    <> ") "
                    <> pretty body
                    <> ")"
                )
                `shouldSatisfy` isRight

  prop "parses a definition" $
    forAll (genDef 5) $ \ast ->
      let text = pretty ast
       in parse text `shouldBe` Right ast

  describe "integer literals" $ do
    prop "any integer rendered by pretty parses back to the same Lit" $ \n ->
      parse (pretty (Lit n)) === Right (Lit (n :: Int))

    prop "a '-' separated from digits by spaces does not parse as an integer literal" $
      forAll (getPositive <$> (arbitrary :: Gen (Positive Int))) $ \n ->
        let spaced = "- " <> Text.pack (show n)
         in parse spaced `shouldSatisfy` (/= Right (Lit (negate n)))

  describe "arithmetic operators" $ do
    -- Operator names are valid identifiers: they parse as Sym and round-trip.
    it "parses + as a symbol" $
      parse "+" `shouldBe` Right (Sym "+")

    it "parses - as a symbol" $
      parse "-" `shouldBe` Right (Sym "-")

    it "parses * as a symbol" $
      parse "*" `shouldBe` Right (Sym "*")

    it "parses div as a symbol" $
      parse "div" `shouldBe` Right (Sym "div")

    it "parses mod as a symbol" $
      parse "mod" `shouldBe` Right (Sym "mod")

    it "parses = as a symbol" $
      parse "=" `shouldBe` Right (Sym "=")

    it "parses < as a symbol" $
      parse "<" `shouldBe` Right (Sym "<")

    it "parses neg as a symbol" $
      parse "neg" `shouldBe` Right (Sym "neg")

    -- Full expression parses and round-trips
    it "parses (+ 3 4) and pretty-prints back" $
      let expr = App (Sym "+") (Lit 3) [Lit 4]
       in parse "(+ 3 4)" `shouldBe` Right expr

    it "pretty (+ 3 4) round-trips through parse" $
      let expr = App (Sym "+") (Lit 3) [Lit 4]
       in parse (pretty expr) `shouldBe` Right expr

    -- desugar of an arithmetic expression produces a normalised integer
    it "desugar and normalise (+ 3 4) == Num 7" $
      normalise (desugar (App (Sym "+") (Lit 3) [Lit 4])) `shouldBe` Eval.Num 7

    it "desugar and normalise (* 2 (+ 1 4)) == Num 10" $
      normalise (desugar (App (Sym "*") (Lit 2) [App (Sym "+") (Lit 1) [Lit 4]]))
        `shouldBe` Eval.Num 10

    -- Comparison operators produce Church booleans
    it "desugar and normalise (= 3 3) == Church True" $
      normalise (desugar (App (Sym "=") (Lit 3) [Lit 3]))
        `shouldBe` Eval.Lam "f" (Eval.Lam "t" (Eval.Var "t"))

    it "desugar and normalise (= 3 4) == Church False" $
      normalise (desugar (App (Sym "=") (Lit 3) [Lit 4]))
        `shouldBe` Eval.Lam "f" (Eval.Lam "t" (Eval.Var "f"))

    it "desugar and normalise (< 2 5) == Church True" $
      normalise (desugar (App (Sym "<") (Lit 2) [Lit 5]))
        `shouldBe` Eval.Lam "f" (Eval.Lam "t" (Eval.Var "t"))

    -- A partial application (missing second arg) pretty-prints as an App
    it "pretty (+ 1) is \"(+ 1)\"" $
      pretty (App (Sym "+") (Lit 1) []) `shouldBe` "(+ 1)"

  it "parses multiline program with blank lines" $ do
    let program =
          "(def id (lam (x) x))\n\
          \\n\
          \(def const (lam (x y) x))\n\
          \(def apply (lam (f x) (f x)))\n\
          \(apply id const)"
        expected =
          [ Def "id" (Abs ["x"] (Sym "x"))
          , Def "const" (Abs ["x", "y"] (Sym "x"))
          , Def "apply" (Abs ["f", "x"] (App (Sym "f") (Sym "x") []))
          , App
              (Sym "apply")
              (Sym "id")
              [Sym "const"]
          ]
    parseProgram program `shouldBe` Right expected

spine :: AST -> Int
spine = \case
  App x _ rest -> 1 + spine x + length rest
  _ -> 0

arity :: AST -> Int
arity = \case
  Abs bind body -> length bind + arity body
  _ -> 0

newtype Identifier = Identifier {unIdent :: Text}
  deriving (Eq, Show)

instance Arbitrary Identifier where
  arbitrary =
    ( Identifier . pack
        <$> ((:) <$> elements initialChars <*> listOf (elements restChars))
    )
      `suchThat` (not . isIntLiteral . unIdent)
   where
    -- An identifier whose first characters would be consumed as an integer
    -- literal by the parser must be excluded: the parser tries `lit` before
    -- `sym`, so any string starting with digits, or with '-' immediately
    -- followed by digits, would produce a Lit rather than a Sym.
    isIntLiteral t =
      case Text.unpack t of
        ('-' : d : _) -> isDigit d
        (d : _) -> isDigit d
        [] -> False
    isDigit c = c >= '0' && c <= '9'

instance Arbitrary AST where
  arbitrary = genAst 10

  shrink = \case
    Sym _ -> []
    Lit _ -> []
    Abs bind body -> case shrink (Identifier <$> bind) of
      [] -> shrink body
      bind' ->
        [ shrunkAbs (unIdent <$> bound) body'
        | bound <- bind'
        , body' <- shrink body
        ]
       where
        shrunkAbs [] b = b
        shrunkAbs bound b = Abs bound b
    App x y [] -> [x, y]
    App x y rest -> (App x y <$> shrink rest) <> [x, y]
    Def _name body -> shrink body -- never generated

genAst :: Int -> Gen AST
genAst 0 = oneof [Sym . unIdent <$> arbitrary, Lit <$> arbitrary]
genAst depth =
  oneof
    [ Sym . unIdent <$> arbitrary
    , Lit <$> arbitrary
    , genAbstraction depth
    , genApplication depth
    ]

genDef :: Int -> Gen AST
genDef n = Def . unIdent <$> arbitrary <*> genAst n

genApplication :: Int -> Gen AST
genApplication depth = do
  App <$> genAst d <*> genAst d <*> reasonablySized (listOf $ genAst d)
 where
  d = depth - 1

genAbstraction :: Int -> Gen AST
genAbstraction depth = do
  x <- unIdent <$> arbitrary
  xs <- fmap unIdent <$> reasonablySized arbitrary
  Abs (x : xs) <$> reasonablySized (genAst $ depth - 1)

reasonablySized :: Gen a -> Gen a
reasonablySized gen = sized $ \size ->
  resize (floor $ sqrt $ fromIntegral @_ @Double size) gen