{-# LANGUAGE LambdaCase #-}
module Minilang.Lambda.EvalSpec (spec) where

import Control.Concurrent (threadDelay)
import Control.Concurrent.Async (race_)
import Data.Aeson (eitherDecodeFileStrict)
import Data.Text (Text, pack)
import qualified Data.Text as Text
import Minilang.API (Challenge (..))
import Minilang.Lambda (loadProgram)
import qualified Minilang.Lambda.DeBruijn as DB
import Minilang.Lambda.DeBruijn (toDeBruijn)
import Minilang.Lambda.Eval (Env, Term (..), Value (..), eval, normalise, builtinEnvValue, newEnv)
import Minilang.Lambda.Infer (newContext)
import Minilang.Lambda.Parser (parseProgram)
import Minilang.REPL.Types (Evaluator (..), Language (..))
import Test.Hspec (Spec, describe, it, runIO, shouldBe)
import Test.Hspec.QuickCheck (prop)
import Test.QuickCheck (NonEmptyList (..), arbitrary, forAll, (===))

spec :: Spec
spec = do
  let y =
        Lam
          "f"
          ( App
              (Lam "x" (App (Var "f") (App (Var "x") (Var "x"))))
              (Lam "x" (App (Var "f") (App (Var "x") (Var "x"))))
          )
      z =
        Lam
          "f"
          ( App
              (Lam "x" (App (Var "f") (Lam "y" (App (App (Var "x") (Var "x")) (Var "y")))))
              (Lam "x" (App (Var "f") (Lam "y" (App (App (Var "x") (Var "x")) (Var "y")))))
          )

      omega =
        App (Lam "x" (App (Var "x") (Var "x"))) (Lam "x" (App (Var "x") (Var "x")))

      fst' = Lam "x" (Lam "y" (Var "x"))
      snd' = Lam "x" (Lam "y" (Var "y"))

  describe "call-by-value eval" $ do
    it "evaluates 'x' as itself" $ do
      let term = Var "x"
      normalise term `shouldBe` Var "x"
    it "evaluates '((λ x. x) y)' as 'y'" $ do
      let term = App (Lam "x" (Var "x")) (Var "y")
      normalise term `shouldBe` Var "y"
    it "evaluates '((λ x y. y) z t y)' as '(t y)'" $ do
      let term = App (App (App snd' (Var "z")) (Var "t")) (Var "y")
      normalise term `shouldBe` App (Var "t") (Var "y")
    it "evaluates '((λ x y. x) z t)' as 'z'" $ do
      let term = App (App fst' (Var "z")) (Var "t")
      normalise term `shouldBe` Var "z"
    it "diverges on evaluating ω" $
      shouldDiverge $
        normalise omega
    it "diverges on evaluating ω applied to a function that use it" $
      shouldDiverge $
        normalise (App fst' omega)
    it "diverges on evaluating ω applied to a function that do not use it" $
      shouldDiverge $
        normalise (App snd' omega)
  describe "Church numerals" $ do
    let zero = Lam "s" (Lam "z" (Var "z"))
        one = Lam "s" (Lam "z" (App (Var "s") (Var "z")))
        two = Lam "s" (Lam "z" (App "s" (App (Var "s") (Var "z"))))
        succ' = Lam "n" (Lam "s" (Lam "z" (App (Var "s") (App (App (Var "n") (Var "s")) (Var "z")))))
        add' = Lam "j" (Lam "k" (Lam "s" (Lam "z" (App (App "j" "s") (App (App "k" "s") "z")))))

    it "normalises successor" $
      normalise (App succ' zero)
        `shouldBe` one

    it "adds two numerals" $
      normalise (App (App add' one) one)
        `shouldBe` two

  describe "integer literals" $ do
    -- An integer literal is already a normal form: normalise must return it unchanged.
    prop "any integer literal normalises to itself" $ \n ->
      normalise (Num n) === Num (n :: Int)

    -- eval on a bare Num must always produce the VNum value form, never a closure
    -- or neutral, regardless of the integer's sign or magnitude.
    prop "eval of any integer literal produces VNum" $ \n ->
      eval [] (Num n) === VNum (n :: Int)

    -- A number is not a function.  Applying it to any argument must leave a stuck
    -- neutral application: App (Num n) (normalise arg).  We test this with free
    -- variables as arguments, which are themselves already in normal form.
    prop "an integer in function position produces a stuck application" $
      forAll arbitrary $ \n ->
        forAll (getNonEmpty <$> arbitrary) $ \x ->
          normalise (App (Num n) (Var (pack x)))
            === App (Num (n :: Int)) (Var (pack x))

  describe "built-in arithmetic" $ do
    -- Basic two-argument reductions
    it "(+ 3 4) normalises to 7" $
      normalise (App (App (Var "+") (Num 3)) (Num 4)) `shouldBe` Num 7

    it "(- 10 3) normalises to 7" $
      normalise (App (App (Var "-") (Num 10)) (Num 3)) `shouldBe` Num 7

    it "(* 2 5) normalises to 10" $
      normalise (App (App (Var "*") (Num 2)) (Num 5)) `shouldBe` Num 10

    it "(div 10 3) normalises to 3" $
      normalise (App (App (Var "div") (Num 10)) (Num 3)) `shouldBe` Num 3

    it "(mod 10 3) normalises to 1" $
      normalise (App (App (Var "mod") (Num 10)) (Num 3)) `shouldBe` Num 1

    it "(= 3 3) normalises to Church True" $
      normalise (App (App (Var "=") (Num 3)) (Num 3))
        `shouldBe` Lam "f" (Lam "t" (Var "t"))

    it "(= 3 4) normalises to Church False" $
      normalise (App (App (Var "=") (Num 3)) (Num 4))
        `shouldBe` Lam "f" (Lam "t" (Var "f"))

    it "(< 2 5) normalises to Church True" $
      normalise (App (App (Var "<") (Num 2)) (Num 5))
        `shouldBe` Lam "f" (Lam "t" (Var "t"))

    it "(< 5 2) normalises to Church False" $
      normalise (App (App (Var "<") (Num 5)) (Num 2))
        `shouldBe` Lam "f" (Lam "t" (Var "f"))

    -- Compose with Church ite: (ite (= 3 3) (Num 1) (Num 2)) == 1
    -- ite c t e = c e t  (True picks second arg, False picks first)
    it "(= 3 3) composes with Church ite to select then-branch" $
      let ite = Lam "c" (Lam "t" (Lam "e" (App (App (Var "c") (Var "e")) (Var "t"))))
      in normalise (App (App (App ite (App (App (Var "=") (Num 3)) (Num 3))) (Num 1)) (Num 2))
           `shouldBe` Num 1

    it "(< 5 2) composes with Church ite to select else-branch" $
      let ite = Lam "c" (Lam "t" (Lam "e" (App (App (Var "c") (Var "e")) (Var "t"))))
      in normalise (App (App (App ite (App (App (Var "<") (Num 5)) (Num 2))) (Num 1)) (Num 2))
           `shouldBe` Num 2

    it "(neg 7) normalises to -7" $
      normalise (App (Var "neg") (Num 7)) `shouldBe` Num (-7)

    it "(neg -3) normalises to 3" $
      normalise (App (Var "neg") (Num (-3))) `shouldBe` Num 3

    -- Nested arithmetic
    it "(* 2 (+ 1 4)) normalises to 10" $
      normalise (App (App (Var "*") (Num 2)) (App (App (Var "+") (Num 1)) (Num 4)))
        `shouldBe` Num 10

    -- Properties
    prop "addition is commutative" $ \x y ->
      normalise (App (App (Var "+") (Num x)) (Num y))
        === normalise (App (App (Var "+") (Num y)) (Num (x :: Int)))

    prop "subtraction is the inverse of addition" $ \x y ->
      normalise (App (App (Var "-") (App (App (Var "+") (Num x)) (Num y))) (Num y))
        === Num (x :: Int)

    prop "multiplication by 0 is 0" $ \x ->
      normalise (App (App (Var "*") (Num x)) (Num 0))
        === Num (0 :: Int)

    prop "neg is its own inverse" $ \x ->
      normalise (App (Var "neg") (App (Var "neg") (Num x)))
        === Num (x :: Int)

    -- Partial application
    it "(+ 1) evaluates to a Builtin waiting for one more arg" $
      eval builtinEnvValue (App (Var "+") (Num 1))
        `shouldBe` Builtin "+" 1 [VNum 1]

    it "partial builtin applied to its last arg reduces" $
      eval builtinEnvValue (App (App (Var "+") (Num 1)) (Num 2))
        `shouldBe` VNum 3

    it "partial builtin passed through a lambda reduces when applied" $
      normalise (App (Lam "f" (App (App (Var "f") (Num 3)) (Num 4))) (Var "+"))
        `shouldBe` Num 7

    -- Stuck / neutral cases
    it "(+ x 1) with free x stays stuck" $
      normalise (App (App (Var "+") (Var "x")) (Num 1))
        `shouldBe` App (App (Var "+") (Var "x")) (Num 1)

    it "(div 1 0) stays stuck" $
      normalise (App (App (Var "div") (Num 1)) (Num 0))
        `shouldBe` App (App (Var "div") (Num 1)) (Num 0)

    it "(mod 5 0) stays stuck" $
      normalise (App (App (Var "mod") (Num 5)) (Num 0))
        `shouldBe` App (App (Var "mod") (Num 5)) (Num 0)

  describe "evaluator equivalence on challenges" $ do
    challenges <- runIO $
      eitherDecodeFileStrict "challenges.json" >>= \case
        Left err -> fail $ "Could not load challenges.json: " <> err
        Right cs -> pure (cs :: [Challenge])

    mapM_ (challengeEquivalenceTest newEnv) challenges

  -- describe "call-by-name eval" $ do
  --   it "evaluates 'x' as itself" $ do
  --     normalise' (Var "x") `shouldBe` Var "x"
  --   it "evaluates '(λ x y . y) ⊥ z' as 'z'" $ do
  --     normalise' (App (App snd' omega) "z") `shouldBe` Var "z"
  --   it "diverges on evaluating ω" $
  --     shouldDiverge $
  --       normalise' omega

-- | For a single challenge expression, parse it and run both evaluators,
-- asserting they produce alpha-equivalent results.
-- Alpha-equivalence is checked by converting both results to De Bruijn-indexed
-- form (which strips variable names) before comparing.
challengeEquivalenceTest :: Env -> Challenge -> Spec
challengeEquivalenceTest env Challenge{level, expr} =
  it ("level " <> show level <> ": " <> Text.unpack (Text.take 60 expr)) $
    case parseProgram expr of
      Left err ->
        fail $ "Parse error for level " <> show level <> ": " <> show err
      Right asts ->
        let runWith ev = loadProgram ev Untyped env newContext asts
            extractTerms (Right (terms, _, _)) = map fst terms
            extractTerms (Left err)            = error $ "Type error: " <> show err
            -- Convert to annotation-erased De Bruijn to compare up to
            -- alpha-equivalence.  Both evaluators may choose different fresh
            -- names for lambda binders (NbE normalises under binders and must
            -- avoid capture); erasing names leaves only structure and indices.
            toAlpha = map (eraseNames . toDeBruijn) . extractTerms
            debruijnTerms = toAlpha (runWith DeBruijn)
            nbeTerms      = toAlpha (runWith NbE)
        in nbeTerms `shouldBe` debruijnTerms

-- | Erase all name annotations from a 'DeBruijn' term, leaving only the
-- structural skeleton (indices, numeric literals, free-variable names, and
-- builtin names).  Two terms that are identical up to alpha-renaming of bound
-- variables will be equal after erasure.
eraseNames :: DB.DeBruijn Text -> DB.DeBruijn ()
eraseNames (DB.FV x)       = DB.FV x
eraseNames (DB.FNum n)     = DB.FNum n
eraseNames (DB.FBuiltin b) = DB.FBuiltin b
eraseNames (DB.Var _ i)    = DB.Var () i
eraseNames (DB.Lam _ body) = DB.Lam () (eraseNames body)
eraseNames (DB.App f a)    = DB.App (eraseNames f) (eraseNames a)

shouldDiverge :: Term -> IO ()
shouldDiverge t =
  race_
    (t `seq` fail "Did not diverge")
    (threadDelay 1000000)