module EvaluationSpec (spec) where

import Data.Set (Set, empty, fromList)
import Evaluation (beta, bind, freeIn, freeVars, freshName, subst)
import Parser (Expr (Abstraction, Application, Variable), parse)
import Test.Hspec

spec :: Spec
spec = do
  describe "freeIn" $ do
    it "variable matches name" $
      "x" `freeIn` Variable "x" `shouldBe` True
    it "variable does not match name" $
      "y" `freeIn` Variable "x" `shouldBe` False
    it "free in abstraction body with different binder" $
      "x" `freeIn` Abstraction "y" (Variable "x") `shouldBe` True
    it "not free when abstraction binder shadows it" $
      "x" `freeIn` Abstraction "x" (Variable "x") `shouldBe` False
    it "outer binder shields entire subtree" $
      "x" `freeIn` Abstraction "x" (Abstraction "y" (Variable "x")) `shouldBe` False
    it "free through nested abstractions when no binder shadows it" $
      "x" `freeIn` Abstraction "y" (Abstraction "z" (Variable "x")) `shouldBe` True
    it "free in left side of application" $
      "x" `freeIn` Application (Variable "x") (Variable "y") `shouldBe` True
    it "free in right side of application" $
      "x" `freeIn` Application (Variable "y") (Variable "x") `shouldBe` True
    it "not free in either side of application" $
      "x" `freeIn` Application (Variable "y") (Variable "z") `shouldBe` False
    it "free on right side even when bound on left side" $
      -- left: λx. x  (x is bound); right: x (x is free)
      "x" `freeIn` Application (Abstraction "x" (Variable "x")) (Variable "x") `shouldBe` True

  describe "freeVars" $ do
    it "single variable" $
      freeVars (Variable "x") `shouldBe` fromList ["x"]
    it "abstraction binds its variable" $
      freeVars (Abstraction "x" (Variable "x")) `shouldBe` empty
    it "abstraction with free variable in body" $
      freeVars (Abstraction "x" (Variable "y")) `shouldBe` fromList ["y"]
    it "application collects from both sides" $
      freeVars (Application (Variable "x") (Variable "y")) `shouldBe` fromList ["x", "y"]
    it "application deduplicates the same variable on both sides" $
      freeVars (Application (Variable "x") (Variable "x")) `shouldBe` fromList ["x"]
    it "abstraction removes its binder from body's free vars" $
      freeVars (Abstraction "x" (Application (Variable "x") (Variable "y"))) `shouldBe` fromList ["y"]
    it "only the unbound occurrence is free across an application" $
      -- left: λx. x (x bound); right: x (x free)
      freeVars (Application (Abstraction "x" (Variable "x")) (Variable "x")) `shouldBe` fromList ["x"]

  describe "freshName" $ do
    it "returns name unchanged when not in the set" $
      freshName "x" empty `shouldBe` "x"
    it "returns name unchanged when other names are in the set" $
      freshName "x" (fromList ["y", "z"]) `shouldBe` "x"
    it "appends one apostrophe when name is taken" $
      freshName "x" (fromList ["x"]) `shouldBe` "x'"
    it "appends two apostrophes when name and name' are both taken" $
      freshName "x" (fromList ["x", "x'"]) `shouldBe` "x''"
    it "appends three apostrophes when name, name', and name'' are all taken" $
      freshName "x" (fromList ["x", "x'", "x''"]) `shouldBe` "x'''"

  describe "subst" $ do
    it "cannot substitute mismatched variables" $
      subst (Variable "x") "y" (Variable "z") `shouldBe` Variable "x"
    it "can substitute matched variables" $
      subst (Variable "x") "x" (Variable "z") `shouldBe` Variable "z"
    it "can substitute nested variables" $
      subst absWithZ "z" absI `shouldBe` absWithZMaker absI
    -- substitution in Application nodes
    it "substitutes in both sides of an application" $
      subst (Application (Variable "x") (Variable "x")) "x" (Variable "z")
        `shouldBe` Application (Variable "z") (Variable "z")
    it "substitutes when content is an abstraction" $
      subst (Variable "x") "x" absI `shouldBe` absI
    it "substitutes complex content into both sides of application" $
      subst (Application (Variable "x") (Variable "x")) "x" absI
        `shouldBe` Application absI absI
    -- bound variable must shadow the substitution
    it "does not substitute a variable that is bound in an abstraction (same binder name)" $
      -- subst (λx. x) "x" z  should be  λx. x   (x is bound, not free)
      subst (Abstraction "x" (Variable "x")) "x" (Variable "z")
        `shouldBe` Abstraction "x" (Variable "x")
    it "does not substitute free variable in body when binder shadows it" $
      -- subst (λx. x y) "x" z  should be  λx. x y  (the x in the body is bound)
      subst (Abstraction "x" (Application (Variable "x") (Variable "y"))) "x" (Variable "z")
        `shouldBe` Abstraction "x" (Application (Variable "x") (Variable "y"))
    -- variable capture: replacement's free variable must not be captured by a binder
    it "does not capture free variables in the replacement (variable capture)" $
      -- subst (λz. x) "x" z  must rename the binder: result is  λz'. z  (or similar)
      -- The replacement 'z' is free, but the abstraction binds 'z' — naive subst captures it.
      -- Expected: the inner z stays free (some fresh binder wraps it, not 'z').
      -- We test this by asserting the result is NOT λz. z (which would be a capture).
      subst (Abstraction "z" (Variable "x")) "x" (Variable "z")
        `shouldNotBe` Abstraction "z" (Variable "z")
  describe "bind" $ do
    it "bind on non-abstraction returns Nothing" $
      bind (Variable "x") (Variable "y") `shouldBe` Nothing
    it "bind identity lambda with argument returns the argument" $
      bind absI (Variable "y") `shouldBe` Just (Variable "y")
    it "bind constant lambda returns the constant" $
      -- (λx. z) applied to y  →  z
      bind (Abstraction "x" (Variable "z")) (Variable "y") `shouldBe` Just (Variable "z")
    it "bind replaces all free occurrences of the bound variable" $
      -- (λx. x x) applied to y  →  y y
      bind selfApp (Variable "y") `shouldBe` Just (Application (Variable "y") (Variable "y"))
    it "bind with abstraction argument substitutes correctly" $
      -- (λx. x) applied to (λi. i)  →  (λi. i)
      bind absI absI `shouldBe` Just absI
  describe "beta reduction" $ do
    -- Normal forms (no redex)
    it "returns Nothing for a bare variable" $
      beta (Variable "x") `shouldBe` Nothing
    it "returns Nothing for a bare abstraction" $
      beta absI `shouldBe` Nothing
    it "returns Nothing when application has no lambda on left" $
      beta (Application (Variable "x") (Variable "y")) `shouldBe` Nothing
    it "returns Nothing with beta normal expression of great complexity" $
      case parse "λa. λb. λc. λd. a (b (c d)) (c (a d b)) (b (a (c d) b) (c d))" of
        Right exp -> beta exp `shouldBe` Nothing
        Left err -> expectationFailure err

    -- Basic reductions
    it "reduces identity: (λi. i) y → y" $
      beta (Application absI (Variable "y"))
        `shouldBe` Just (Variable "y")
    it "reduces constant function: (λx. z) y → z" $
      beta (Application (Abstraction "x" (Variable "z")) (Variable "y"))
        `shouldBe` Just (Variable "z")

    -- Self-application
    it "(λx. x x)(λi. i) → (λi. i)(λi. i)" $
      beta (Application selfApp absI)
        `shouldBe` Just (Application absI absI)
    it "(λi. i)(λi. i) → (λi. i)" $
      beta (Application absI absI)
        `shouldBe` Just absI

    -- Redex not at top level
    it "reduces redex in function position: ((λi. i) y) z → y z" $
      beta (Application (Application absI (Variable "y")) (Variable "z"))
        `shouldBe` Just (Application (Variable "y") (Variable "z"))
    it "reduces redex inside lambda body: λz. (λi. i) y → λz. y" $
      beta (Abstraction "z" (Application absI (Variable "y")))
        `shouldBe` Just (Abstraction "z" (Variable "y"))

    -- Redex in argument position (function side is already normal)
    it "reduces redex in argument position when function is a variable: x ((λi. i) y) → x y" $
      beta (Application (Variable "x") (Application absI (Variable "y")))
        `shouldBe` Just (Application (Variable "x") (Variable "y"))
    it "reduces redex in argument position when function is a normal abstraction: (λa. a) ((λi. i) y) reduces function first" $
      -- leftmost-outermost: the outer application is the redex, fires first
      beta (Application absI (Application absI (Variable "y")))
        `shouldBe` Just (Application absI (Variable "y"))

    -- Omega combinator (divergent): (λx. x x)(λx. x x) → (λx. x x)(λx. x x)
    it "omega combinator steps back to itself" $
      beta (Application omega omega) `shouldBe` Just (Application omega omega)

    -- Multi-step: applying beta repeatedly converges
    it "iterating beta on (λi. i)((λi. i) y) reaches y in 2 steps" $
      -- step 1: (λi. i)((λi. i) y)  → (λi. i) y   [outer redex fires first]
      -- step 2: (λi. i) y            → y
      let step1 = beta (Application absI (Application absI (Variable "y")))
          step2 = step1 >>= beta
       in step2 `shouldBe` Just (Variable "y")

    -- Substitution must not be performed under a shadowing binder
    it "beta does not substitute into a shadowed variable: (λx. λx. x) y → λx. x" $
      -- The inner λx re-binds x, so the outer x should not be substituted into it
      beta (Application (Abstraction "x" (Abstraction "x" (Variable "x"))) (Variable "y"))
        `shouldBe` Just (Abstraction "x" (Variable "x"))
  where
    absWithZMaker z = Abstraction "x" (Application z (Variable "x"))
    absWithZ = absWithZMaker (Variable "z")
    absI    = Abstraction "i" (Variable "i")
    selfApp = Abstraction "x" (Application (Variable "x") (Variable "x"))
    omega   = Abstraction "x" (Application (Variable "x") (Variable "x"))