{- |
Module      : Llama.Split
Description : High level Split interface for llama-cpp
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
-}
module Llama.Split (
    splitPath
    , splitPrefix
   , printSystemInfo
) where

import Llama.Internal.Foreign
import Foreign
import Foreign.C.String

-- | Split a path into a prefix and a split number
splitPath :: FilePath -> Int -> Int -> IO (Either String String)
splitPath :: FilePath -> Int -> Int -> IO (Either FilePath FilePath)
splitPath FilePath
pathPrefix Int
splitNo Int
splitCount = do
  FilePath
-> (CString -> IO (Either FilePath FilePath))
-> IO (Either FilePath FilePath)
forall a. FilePath -> (CString -> IO a) -> IO a
withCString FilePath
pathPrefix ((CString -> IO (Either FilePath FilePath))
 -> IO (Either FilePath FilePath))
-> (CString -> IO (Either FilePath FilePath))
-> IO (Either FilePath FilePath)
forall a b. (a -> b) -> a -> b
$ \CString
cPathPrefix -> do
    Int
-> (CString -> IO (Either FilePath FilePath))
-> IO (Either FilePath FilePath)
forall a b. Storable a => Int -> (Ptr a -> IO b) -> IO b
allocaArray Int
256 ((CString -> IO (Either FilePath FilePath))
 -> IO (Either FilePath FilePath))
-> (CString -> IO (Either FilePath FilePath))
-> IO (Either FilePath FilePath)
forall a b. (a -> b) -> a -> b
$ \CString
splitPathPtr -> do
      requiredLen <- CString -> CSize -> CString -> CInt -> CInt -> IO CInt
c_llama_split_path CString
splitPathPtr CSize
256 CString
cPathPrefix (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
splitNo) (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
splitCount)
      if requiredLen < 0
        then return $ Left "Failed to split path"
        else do
          str <- peekCString splitPathPtr
          return $ Right str

-- | Get the prefix from a split path
splitPrefix :: FilePath -> Int -> Int -> IO (Either String String)
splitPrefix :: FilePath -> Int -> Int -> IO (Either FilePath FilePath)
splitPrefix FilePath
splitPath_ Int
splitNo Int
splitCount = do
  FilePath
-> (CString -> IO (Either FilePath FilePath))
-> IO (Either FilePath FilePath)
forall a. FilePath -> (CString -> IO a) -> IO a
withCString FilePath
splitPath_ ((CString -> IO (Either FilePath FilePath))
 -> IO (Either FilePath FilePath))
-> (CString -> IO (Either FilePath FilePath))
-> IO (Either FilePath FilePath)
forall a b. (a -> b) -> a -> b
$ \CString
cSplitPath -> do
    Int
-> (CString -> IO (Either FilePath FilePath))
-> IO (Either FilePath FilePath)
forall a b. Storable a => Int -> (Ptr a -> IO b) -> IO b
allocaArray Int
256 ((CString -> IO (Either FilePath FilePath))
 -> IO (Either FilePath FilePath))
-> (CString -> IO (Either FilePath FilePath))
-> IO (Either FilePath FilePath)
forall a b. (a -> b) -> a -> b
$ \CString
splitPrefixPtr -> do
      requiredLen <- CString -> CSize -> CString -> CInt -> CInt -> IO CInt
c_llama_split_prefix CString
splitPrefixPtr CSize
256 CString
cSplitPath (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
splitNo) (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
splitCount)
      if requiredLen < 0
        then return $ Left "Failed to get prefix from split path"
        else do
          str <- peekCString splitPrefixPtr
          return $ Right str

-- | Print system information
printSystemInfo :: IO String
printSystemInfo :: IO FilePath
printSystemInfo = do
  systemInfo <- IO CString
c_llama_print_system_info
  peekCString systemInfo