{-# 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, escapeString)
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)"

  describe "string literals" $ do
    -- Round-trip: parse . pretty == id for all safe StrLit values
    prop "any StrLit rendered by pretty parses back to the same StrLit" $
      forAll genStrLit $ \ast ->
        parse (pretty ast) === Right ast

    it "parses empty string literal" $
      parse "\"\"" `shouldBe` Right (StrLit "")

    it "parses simple string literal" $
      parse "\"hello\"" `shouldBe` Right (StrLit "hello")

    it "parses string with spaces" $
      parse "\"foo bar\"" `shouldBe` Right (StrLit "foo bar")

    it "parses escaped double-quote inside string" $
      parse "\"say \\\"hi\\\"\"" `shouldBe` Right (StrLit "say \"hi\"")

    it "parses escaped backslash inside string" $
      parse "\"back\\\\slash\"" `shouldBe` Right (StrLit "back\\slash")

    it "pretty-prints string with double-quote escaped" $
      pretty (StrLit "say \"hi\"") `shouldBe` "\"say \\\"hi\\\"\""

    it "pretty-prints string with backslash escaped" $
      pretty (StrLit "back\\slash") `shouldBe` "\"back\\\\slash\""

    -- desugar / normalise
    it "desugar and normalise string self-evaluates" $
      normalise (desugar (StrLit "hello")) `shouldBe` Eval.Str "hello"

    it "desugar and normalise str-concat" $
      normalise (desugar (App (Sym "str-concat") (StrLit "foo") [StrLit "bar"]))
        `shouldBe` Eval.Str "foobar"

    it "desugar and normalise str-length" $
      normalise (desugar (App (Sym "str-length") (StrLit "hello") []))
        `shouldBe` Eval.Num 5

    it "desugar and normalise str-slice" $
      normalise (desugar (App (Sym "str-slice") (Lit 1) [Lit 3, StrLit "abcde"]))
        `shouldBe` Eval.Str "bc"

    it "desugar and normalise str= equal strings gives Church True" $
      normalise (desugar (App (Sym "str=") (StrLit "abc") [StrLit "abc"]))
        `shouldBe` Eval.Lam "f" (Eval.Lam "t" (Eval.Var "t"))

    it "desugar and normalise str= unequal strings gives Church False" $
      normalise (desugar (App (Sym "str=") (StrLit "abc") [StrLit "xyz"]))
        `shouldBe` Eval.Lam "f" (Eval.Lam "t" (Eval.Var "f"))

    it "partial str-concat pretty-prints as application" $
      pretty (App (Sym "str-concat") (StrLit "hello") [])
        `shouldBe` "(str-concat \"hello\")"

  describe "strings containing parentheses" $ do
    -- The parser must not treat ( or ) inside a string literal as
    -- delimiters for an application or lambda expression.

    it "parses a string containing only an opening paren" $
      parse "\"(\"" `shouldBe` Right (StrLit "(")

    it "parses a string containing only a closing paren" $
      parse "\")\"" `shouldBe` Right (StrLit ")")

    it "parses a string containing balanced parens" $
      parse "\"(hello)\"" `shouldBe` Right (StrLit "(hello)")

    it "parses a string containing doubly-nested parens" $
      parse "\"((nested))\"" `shouldBe` Right (StrLit "((nested))")

    it "parses a string containing unbalanced parens" $
      parse "\"a(b\"" `shouldBe` Right (StrLit "a(b")

    it "parses a string that looks like a lambda expression" $
      parse "\"(lam (x) x)\"" `shouldBe` Right (StrLit "(lam (x) x)")

    it "parses a string that looks like an application" $
      parse "\"(f x)\"" `shouldBe` Right (StrLit "(f x)")

    -- Round-trip: pretty-printing a StrLit with parens re-parses correctly
    it "pretty-prints and re-parses string with opening paren" $
      parse (pretty (StrLit "(")) `shouldBe` Right (StrLit "(")

    it "pretty-prints and re-parses string with balanced parens" $
      parse (pretty (StrLit "(hello)")) `shouldBe` Right (StrLit "(hello)")

    it "pretty-prints and re-parses string with doubly-nested parens" $
      parse (pretty (StrLit "((x))")) `shouldBe` Right (StrLit "((x))")

    -- Property: any string containing only parens and letters round-trips
    prop "strings with parentheses round-trip through pretty and parse" $
      forAll (StrLit . Text.pack <$> listOf (elements "abcdef()")) $ \ast ->
        parse (pretty ast) === Right ast

    -- The paren-containing string is correctly passed through an application
    it "paren string survives beta-reduction" $
      normalise (desugar (App (Abs ["x"] (Sym "x")) (StrLit "(hello)") []))
        `shouldBe` Eval.Str "(hello)"

    -- str-concat of two paren strings
    it "str-concat of paren strings produces correct result" $
      normalise (desugar (App (Sym "str-concat") (StrLit "(") [StrLit ")"]))
        `shouldBe` Eval.Str "()"

    it "str-length of paren string counts parens" $
      normalise (desugar (App (Sym "str-length") (StrLit "(hello)") []))
        `shouldBe` Eval.Num 7

  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 _ -> []
    StrLit _ -> []
    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, genStrLit]
genAst depth =
  oneof
    [ Sym . unIdent <$> arbitrary
    , Lit <$> arbitrary
    , genStrLit
    , genAbstraction depth
    , genApplication depth
    ]

-- | Generate a random string literal. We restrict the alphabet to printable
-- ASCII excluding @\"@ and @\\@ to keep round-trip tests simple; the escape
-- logic is tested separately.
genStrLit :: Gen AST
genStrLit = StrLit . pack <$> listOf (elements safeChars)
  where
    safeChars = filter (\c -> c /= '"' && c /= '\\') [' '..'~']

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