{-# LANGUAGE OverloadedStrings #-}
module Main (main) where

import Data.Foldable (forM_)
import Embeddings (chunkArticle, tokensApprox, chunkMaxTokens, Embedding(..))
import Test.Hspec
import qualified Data.Text as T

main :: IO ()
main = hspec $ do
  describe "Chunking" $ do
    it "Splits a long article into multiple chunks" testChunking
    it "Retains a short article as one chunk" testShortArticleOneChunk
    it "Prefixes each chunk with the title" testTitlePrefix
    it "Chunks to small enough items" testTokenLimit


testChunking :: IO ()
testChunking = do
  let longHtml = T.replicate 1000 "<p>This is a paragraph with some text content that should be long enough to create multiple chunks when repeated many times.</p>"
      chunks = chunkArticle "Test Article" longHtml
      n = length chunks
  n `shouldBe` 7

testShortArticleOneChunk :: IO ()
testShortArticleOneChunk = do
  let shortHtml = "<p>A short paragraph.</p>"
      chunks = chunkArticle "Short" shortHtml
  length chunks `shouldBe` 1

testTitlePrefix :: IO ()
testTitlePrefix = do
  let chunks = chunkArticle "My Title" "<p>Some content.</p>"
  forM_ chunks $ \chunk ->
    chunk `shouldSatisfy` ("Title: My Title\n\n" `T.isPrefixOf`)

testTokenLimit :: IO ()
testTokenLimit = do
  let longHtml = T.replicate 1000 "<p>This is a paragraph with some text content that should be long enough to create multiple chunks when repeated many times.</p>"
      chunks = chunkArticle "Test Article" longHtml
      counts = map tokensApprox chunks
  counts `shouldSatisfy` all (<= chunkMaxTokens)