content: {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeFamilies #-}
  
  module Test.Langchain.Runnable.Chains (tests) where
  
  import Langchain.Error (LangchainError, llmError)
  import Langchain.Runnable.Chain
  import Langchain.Runnable.Core
  import Test.Tasty (TestTree, testGroup)
  import Test.Tasty.HUnit (assertEqual, testCase)
  
  addOne :: MockRunnable Int Int
  addOne = MockRunnable (\x -> return $ Right (x + 1))
  
  multiplyByTwo :: MockRunnable Int Int
  multiplyByTwo = MockRunnable (\x -> return $ Right (x * 2))
  
  evenCheck :: MockRunnable Int Bool
  evenCheck = MockRunnable $ return . Right . even
  
  failingMock :: MockRunnable a b
  failingMock = MockRunnable (\_ -> return $ Left (llmError "Mock error" Nothing Nothing))
  
  newtype MockRunnable a b = MockRunnable {runMock :: a -> IO (Either LangchainError b)}
  
  instance Runnable (MockRunnable a b) where
    type RunnableInput (MockRunnable a b) = a
    type RunnableOutput (MockRunnable a b) = b
    invoke = runMock
  
  tests :: TestTree
  tests =
    testGroup
      "Runnable Chain Tests"
      [ testGroup
          "RunnableBranch Tests"
          [ testCase "Selects first matching branch" $ do
              let branch1 =
                    RunnableBranch
                      [ ((== 1), addOne)
                      , ((== 2), multiplyByTwo)
                      ]
                      failingMock
              result <- runBranch branch1 1
              assertEqual "Should choose addOne branch" (Right 2) result
          , testCase "Uses default when no conditions match" $ do
              let defaultBranch = RunnableBranch [] addOne
              result <- runBranch defaultBranch 5
              assertEqual "Should use default" (Right 6) result
          ]
      , testGroup
          "RunnableMap Tests"
          [ testCase "Applies input/output transformations" $ do
              let inputMap = (* 2)
                  outputMap = (+ 1)
                  mapped = RunnableMap inputMap outputMap addOne
              result <- runMap mapped 3 -- 3*2=6 → addOne →7 → +1 →8
              assertEqual "Transformations applied" (Right 8) result
          ]
      , testGroup
          "RunnableSequence Tests"
          [ testCase "Executes sequence in order" $ do
              let sequence0 = buildSequence addOne multiplyByTwo
              result <- runSequence sequence0 2 -- 2+1=3 → *2=6
              assertEqual "Sequence executed" (Right 6) result
  
              {-
              , testCase "Handles multi-step sequences" $ do
                  let sequence_ = (addOne |>> multiplyByTwo) |>> evenCheck
                  result <- sequence_ 3 -- 3+1=4 → *2=8 → even → True
                  assertEqual "Three-step sequence" (Right True) result
                  -}
          ]
      , testGroup
          "Chain Operator Tests"
          [ testCase "Chains two runnables" $ do
              let pipeline = addOne |>> multiplyByTwo
              result <- pipeline 3
              assertEqual "3+1=4 → *2=8" (Right 8) result
          , testCase "Propagates errors in chain" $ do
              let pipeline = failingMock |>> multiplyByTwo
              result <- pipeline ()
              assertEqual "Error in first step" (Left (llmError "Mock error" Nothing Nothing)) result
          ]
      , testGroup
          "Branch Tests"
          [ testCase "Runs parallel branches" $ do
              result <- branch evenCheck addOne 4
              assertEqual "Both branches run" (Right (True, 5)) result
          , testCase "Handles branch errors" $ do
              result <- branch failingMock addOne 5
              assertEqual
                "Left error in first branch"
                (Left (llmError "Mock error" Nothing Nothing) :: Either LangchainError (Bool, Int))
                result
          ]
      ]
  