module EvaluationSpec (spec) where

import Data.Set (Set, empty, fromList)
import Evaluation (alphaEquiv, beta, bind, freeIn, freeVars, freshName, run, 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 "alphaEquiv" $ do

    -- Reflexivity
    it "a variable is alpha-equivalent to itself" $
      alphaEquiv (Variable "x") (Variable "x") `shouldBe` True
    it "an abstraction is alpha-equivalent to itself" $
      alphaEquiv (Abstraction "x" (Variable "x")) (Abstraction "x" (Variable "x")) `shouldBe` True
    it "a nested term is alpha-equivalent to itself" $
      alphaEquiv (Abstraction "x" (Abstraction "y" (Application (Variable "x") (Variable "y"))))
                 (Abstraction "x" (Abstraction "y" (Application (Variable "x") (Variable "y"))))
        `shouldBe` True

    -- Renaming bound variables (True)
    it "identity with different binder names: λx. x ≡ λy. y" $
      alphaEquiv (Abstraction "x" (Variable "x"))
                 (Abstraction "y" (Variable "y"))
        `shouldBe` True
    it "constant body with different binder names: λx. z ≡ λy. z" $
      alphaEquiv (Abstraction "x" (Variable "z"))
                 (Abstraction "y" (Variable "z"))
        `shouldBe` True
    it "nested: λx. λy. x ≡ λa. λb. a (outer reference)" $
      alphaEquiv (Abstraction "x" (Abstraction "y" (Variable "x")))
                 (Abstraction "a" (Abstraction "b" (Variable "a")))
        `shouldBe` True
    it "nested: λx. λy. y ≡ λa. λb. b (inner reference)" $
      alphaEquiv (Abstraction "x" (Abstraction "y" (Variable "y")))
                 (Abstraction "a" (Abstraction "b" (Variable "b")))
        `shouldBe` True
    it "nested application body: λx. λy. x y ≡ λa. λb. a b" $
      alphaEquiv (Abstraction "x" (Abstraction "y" (Application (Variable "x") (Variable "y"))))
                 (Abstraction "a" (Abstraction "b" (Application (Variable "a") (Variable "b"))))
        `shouldBe` True
    it "self-application body: λx. x x ≡ λy. y y" $
      alphaEquiv (Abstraction "x" (Application (Variable "x") (Variable "x")))
                 (Abstraction "y" (Application (Variable "y") (Variable "y")))
        `shouldBe` True

    -- Symmetry
    it "alpha equivalence is symmetric" $
      alphaEquiv (Abstraction "x" (Variable "x"))
                 (Abstraction "y" (Variable "y"))
        `shouldBe`
      alphaEquiv (Abstraction "y" (Variable "y"))
                 (Abstraction "x" (Variable "x"))

    -- Free variables are name-sensitive (False)
    it "distinct free variables are not equivalent" $
      alphaEquiv (Variable "x") (Variable "y") `shouldBe` False
    it "different free variables in abstraction body: λx. y ≢ λx. z" $
      alphaEquiv (Abstraction "x" (Variable "y"))
                 (Abstraction "x" (Variable "z"))
        `shouldBe` False
    it "bound body vs free body: λx. x ≢ λx. y" $
      alphaEquiv (Abstraction "x" (Variable "x"))
                 (Abstraction "x" (Variable "y"))
        `shouldBe` False

    -- Structurally different terms (False)
    it "abstraction vs bare variable are not equivalent" $
      alphaEquiv (Variable "x")
                 (Abstraction "x" (Variable "x"))
        `shouldBe` False
    it "different depths: λx. x ≢ λx. λy. y" $
      alphaEquiv (Abstraction "x" (Variable "x"))
                 (Abstraction "x" (Abstraction "y" (Variable "y")))
        `shouldBe` False
    it "argument order matters in application: x y ≢ y x" $
      alphaEquiv (Application (Variable "x") (Variable "y"))
                 (Application (Variable "y") (Variable "x"))
        `shouldBe` False

    -- Shadowing / name reuse
    it "consistent inner shadowing: λx. λx. x ≡ λy. λy. y" $
      alphaEquiv (Abstraction "x" (Abstraction "x" (Variable "x")))
                 (Abstraction "y" (Abstraction "y" (Variable "y")))
        `shouldBe` True
    it "inner shadow differs from outer reference: λx. λx. x ≢ λx. λy. x" $
      alphaEquiv (Abstraction "x" (Abstraction "x" (Variable "x")))
                 (Abstraction "x" (Abstraction "y" (Variable "x")))
        `shouldBe` False

    -- Free variables unchanged across binder rename
    it "free variable in body is preserved: λx. y ≡ λz. y" $
      alphaEquiv (Abstraction "x" (Variable "y"))
                 (Abstraction "z" (Variable "y"))
        `shouldBe` True
    it "free vs bound distinction preserved: λy. y ≢ λy. z" $
      alphaEquiv (Abstraction "y" (Variable "y"))
                 (Abstraction "y" (Variable "z"))
        `shouldBe` False

  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"))

  describe "run" $ do
    -- Normal forms: run should be a no-op
    it "returns a bare variable unchanged" $
      run (Variable "x") `shouldBe` Variable "x"
    it "returns a bare abstraction unchanged" $
      run absI `shouldBe` absI
    it "returns a normal application (no redex) unchanged" $
      run (Application (Variable "x") (Variable "y")) `shouldBe` Application (Variable "x") (Variable "y")
    it "returns a complex beta-normal term unchanged" $
      case parse "λa. λb. λc. a (b c)" of
        Right expr -> run expr `shouldBe` expr
        Left err -> expectationFailure err

    -- Single-step reductions
    it "reduces identity applied to a variable: (λi. i) y → y" $
      run (Application absI (Variable "y")) `shouldBe` Variable "y"
    it "reduces constant function: (λx. z) y → z" $
      run (Application (Abstraction "x" (Variable "z")) (Variable "y")) `shouldBe` Variable "z"

    -- Multi-step reductions
    it "reduces two nested identity applications to normal form: (λi. i)((λi. i) y) → y" $
      run (Application absI (Application absI (Variable "y"))) `shouldBe` Variable "y"
    it "reduces self-application of identity: (λx. x x)(λi. i) → (λi. i)" $
      run (Application selfApp absI) `shouldBeAlphaEq` absI

    -- Combinator reductions
    it "fully reduces K combinator: K a b → a" $
      run (Application (Application k (Variable "a")) (Variable "b")) `shouldBe` Variable "a"
    it "reduces S K K to identity" $
      run (Application (Application s k) k) `shouldBeAlphaEq` absI

    -- Shadowing
    it "does not substitute into a shadowed binder: (λx. λx. x) y → λx. x" $
      run (Application (Abstraction "x" (Abstraction "x" (Variable "x"))) (Variable "y"))
        `shouldBe` Abstraction "x" (Variable "x")

    -- Capture-avoidance
    it "avoids variable capture: (λx. λz. x) z does not reduce to identity" $
      -- Naive substitution would yield λz. z (identity), capturing the free z.
      -- Correct result renames the binder: λz'. z (a constant-z function).
      run (Application (Abstraction "x" (Abstraction "z" (Variable "x"))) (Variable "z"))
        `shouldBeAlphaEq'` Abstraction "z" (Variable "z")

    -- Idempotency: run (run e) == run e
    it "run is idempotent on a normal form" $
      run (run (Variable "x")) `shouldBe` run (Variable "x")
    it "run is idempotent on a reducible expression" $
      run (run (Application absI (Variable "y"))) `shouldBe` run (Application absI (Variable "y"))

  where
    -- | Assert that two expressions are alpha-equivalent.
    shouldBeAlphaEq :: Expr -> Expr -> Expectation
    shouldBeAlphaEq actual expected =
      alphaEquiv actual expected
        `shouldBe` True
    -- | Assert that two expressions are NOT alpha-equivalent.
    shouldBeAlphaEq' :: Expr -> Expr -> Expectation
    shouldBeAlphaEq' actual expected =
      alphaEquiv actual expected
        `shouldBe` False
    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"))
    k       = Abstraction "x" (Abstraction "y" (Variable "x"))
    s       = Abstraction "x" (Abstraction "y" (Abstraction "z"
                (Application (Application (Variable "x") (Variable "z"))
                             (Application (Variable "y") (Variable "z")))))