{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeApplications #-}

module Minilang.Lambda.ParserSpec where

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, pretty, restChars)
import Test.Hspec (Spec, parallel, shouldBe)
import Test.Hspec.QuickCheck (prop)
import Test.QuickCheck (Arbitrary (..), Gen, NonEmptyList (..), counterexample, elements, forAll, forAllShrink, listOf, oneof, resize, sized, (===), (==>))

-- ((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" $ \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 a let-binding as a function application" $ \(Identifier ident) ->
    forAll (genAst 3) $ \binding ->
      forAll (genAst 2) $ \body ->
        let fun = Abs [ident] body
            app = App fun binding []
         in parse ("(let (" <> ident <> " " <> pretty binding <> ") " <> pretty body <> ")")
              `shouldBe` Right app

spine :: AST -> Int
spine = \case
  App x y 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))

instance Arbitrary AST where
  arbitrary = genAst 10

  shrink = \case
    Sym _ -> []
    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]

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

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