{-# LANGUAGE OverloadedStrings #-}

-- SPDX-FileCopyrightText: Copyright (c) 2025 Objectionary.com
-- SPDX-License-Identifier: MIT

module Functions where

import Ast
import Builder
import Control.Exception (throwIO)
import Control.Monad (replicateM)
import qualified Data.ByteString.Char8 as B
import Data.Char (intToDigit)
import Data.IORef
import Data.Set (Set)
import qualified Data.Set
import GHC.IO (unsafePerformIO)
import Matcher
import Misc
import Numeric (showHex)
import Pretty
import Regexp
import System.Random (randomRIO)
import Term
import Text.Printf (printf)
import Yaml

randomStrings :: IORef (Set String)
{-# NOINLINE randomStrings #-}
randomStrings :: IORef (Set String)
randomStrings = IO (IORef (Set String)) -> IORef (Set String)
forall a. IO a -> a
unsafePerformIO (Set String -> IO (IORef (Set String))
forall a. a -> IO (IORef a)
newIORef Set String
forall a. Set a
Data.Set.empty)

argToStrBytes :: ExtraArgument -> Subst -> Program -> IO String
argToStrBytes :: ExtraArgument -> Subst -> Program -> IO String
argToStrBytes (ArgBytes Bytes
bytes) Subst
subst Program
_ = do
  Bytes
bts <- Bytes -> Subst -> IO Bytes
buildBytesThrows Bytes
bytes Subst
subst
  String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bytes -> String
btsToUnescapedStr Bytes
bts)
argToStrBytes (ArgExpression Expression
expr) Subst
subst Program
prog = do
  (TeBytes Bytes
bts) <- String -> [ExtraArgument] -> Subst -> Program -> IO Term
buildTermFromFunction String
"dataize" [Expression -> ExtraArgument
ArgExpression Expression
expr] Subst
subst Program
prog
  String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bytes -> String
btsToUnescapedStr Bytes
bts)
argToStrBytes ExtraArgument
arg Subst
_ Program
_ = IOError -> IO String
forall e a. Exception e => e -> IO a
throwIO (String -> IOError
userError (String -> String -> String
forall r. PrintfType r => String -> r
printf String
"Can't extract bytes from given argument: %s" (ExtraArgument -> String
prettyExtraArg ExtraArgument
arg)))

buildTermFromFunction :: String -> [ExtraArgument] -> Subst -> Program -> IO Term
buildTermFromFunction :: String -> [ExtraArgument] -> Subst -> Program -> IO Term
buildTermFromFunction String
"contextualize" [ArgExpression Expression
expr, ArgExpression Expression
context] Subst
subst Program
prog = do
  (Expression
expr', Expression
_) <- Expression -> Subst -> IO (Expression, Expression)
buildExpressionThrows Expression
expr Subst
subst
  (Expression
context', Expression
_) <- Expression -> Subst -> IO (Expression, Expression)
buildExpressionThrows Expression
context Subst
subst
  Term -> IO Term
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expression -> Term
TeExpression (Expression -> Expression -> Program -> Expression
contextualize Expression
expr' Expression
context' Program
prog))
buildTermFromFunction String
"contextualize" [ExtraArgument]
_ Subst
_ Program
_ = IOError -> IO Term
forall e a. Exception e => e -> IO a
throwIO (String -> IOError
userError String
"Function contextualize() requires exactly 2 arguments as expression")
buildTermFromFunction String
"scope" [ArgExpression Expression
expr] Subst
subst Program
_ = do
  (Expression
expr', Expression
scope) <- Expression -> Subst -> IO (Expression, Expression)
buildExpressionThrows Expression
expr Subst
subst
  Term -> IO Term
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expression -> Term
TeExpression Expression
scope)
buildTermFromFunction String
"scope" [ExtraArgument]
_ Subst
_ Program
_ = IOError -> IO Term
forall e a. Exception e => e -> IO a
throwIO (String -> IOError
userError String
"Function scope() requires exactly 1 argument as expression")
buildTermFromFunction String
"random-tau" [ExtraArgument]
args Subst
subst Program
_ = do
  [String]
attrs <- [ExtraArgument] -> IO [String]
argsToAttrs [ExtraArgument]
args
  Term -> IO Term
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Attribute -> Term
TeAttribute (String -> Attribute
AtLabel (Integer -> [String] -> String
randomTau Integer
0 [String]
attrs)))
  where
    argsToAttrs :: [ExtraArgument] -> IO [String]
    argsToAttrs :: [ExtraArgument] -> IO [String]
argsToAttrs [] = [String] -> IO [String]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    argsToAttrs (ExtraArgument
arg : [ExtraArgument]
rest) = case ExtraArgument
arg of
      ArgExpression Expression
_ -> [ExtraArgument] -> IO [String]
argsToAttrs [ExtraArgument]
rest
      ArgAttribute Attribute
attr -> do
        Attribute
attr' <- Attribute -> Subst -> IO Attribute
buildAttributeThrows Attribute
attr Subst
subst
        [String]
rest' <- [ExtraArgument] -> IO [String]
argsToAttrs [ExtraArgument]
rest
        [String] -> IO [String]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Attribute -> String
prettyAttribute Attribute
attr' String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
rest')
      ArgBinding Binding
bd -> do
        [Binding]
bds <- Binding -> Subst -> IO [Binding]
buildBindingThrows Binding
bd Subst
subst
        [String]
rest' <- [ExtraArgument] -> IO [String]
argsToAttrs [ExtraArgument]
rest
        [String] -> IO [String]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Binding] -> [String]
attrsFromBindings [Binding]
bds [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ [String]
rest')
      ArgBytes Bytes
_ -> IOError -> IO [String]
forall e a. Exception e => e -> IO a
throwIO (String -> IOError
userError String
"Bytes can't be argument of random-tau() function")
    attrsFromBindings :: [Binding] -> [String]
    attrsFromBindings :: [Binding] -> [String]
attrsFromBindings [] = []
    attrsFromBindings (Binding
bd : [Binding]
bds) =
      let attr :: Attribute
attr = case Binding
bd of
            BiTau Attribute
attr Expression
_ -> Attribute
attr
            BiDelta Bytes
_ -> Attribute
AtDelta
            BiLambda String
_ -> Attribute
AtLambda
            BiVoid Attribute
attr -> Attribute
attr
       in Attribute -> String
prettyAttribute Attribute
attr String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [Binding] -> [String]
attrsFromBindings [Binding]
bds
    randomTau :: Integer -> [String] -> String
    randomTau :: Integer -> [String] -> String
randomTau Integer
idx [String]
attrs =
      let cactoos :: String
cactoos = String
"a🌵"
          tau :: String
tau = if Integer
idx Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 then String
cactoos else String
cactoos String -> String -> String
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
idx
       in if String
tau String -> [String] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [String]
attrs then Integer -> [String] -> String
randomTau (Integer
idx Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1) [String]
attrs else String
tau
buildTermFromFunction String
"dataize" [ArgBytes Bytes
bytes] Subst
subst Program
_ = do
  Bytes
bts <- Bytes -> Subst -> IO Bytes
buildBytesThrows Bytes
bytes Subst
subst
  Term -> IO Term
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bytes -> Term
TeBytes Bytes
bts)
buildTermFromFunction String
"dataize" [ArgExpression Expression
expr] Subst
subst Program
_ = do
  (Expression
expr', Expression
_) <- Expression -> Subst -> IO (Expression, Expression)
buildExpressionThrows Expression
expr Subst
subst
  case Expression
expr' of
    DataObject String
_ Bytes
bytes -> Term -> IO Term
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bytes -> Term
TeBytes Bytes
bytes)
    Expression
_ -> IOError -> IO Term
forall e a. Exception e => e -> IO a
throwIO (String -> IOError
userError String
"Only data objects and bytes are supported by 'dataize' function now")
buildTermFromFunction String
"dataize" [ExtraArgument]
_ Subst
_ Program
_ = IOError -> IO Term
forall e a. Exception e => e -> IO a
throwIO (String -> IOError
userError String
"Function dataize() requires exactly 1 argument as expression")
buildTermFromFunction String
"concat" [ExtraArgument]
args Subst
subst Program
prog = do
  [String]
args' <- (ExtraArgument -> IO String) -> [ExtraArgument] -> IO [String]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (\ExtraArgument
arg -> ExtraArgument -> Subst -> Program -> IO String
argToStrBytes ExtraArgument
arg Subst
subst Program
prog) [ExtraArgument]
args
  Term -> IO Term
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expression -> Term
TeExpression (Bytes -> Expression
DataString (String -> Bytes
strToBts ([String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [String]
args'))))
buildTermFromFunction String
"sed" [ExtraArgument
tgt, ExtraArgument
ptn] Subst
subst Program
prog = do
  [String
tgt', String
ptn'] <- (ExtraArgument -> IO String) -> [ExtraArgument] -> IO [String]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (\ExtraArgument
arg -> ExtraArgument -> Subst -> Program -> IO String
argToStrBytes ExtraArgument
arg Subst
subst Program
prog) [ExtraArgument
tgt, ExtraArgument
ptn]
  (ByteString
pat, ByteString
rep, Bool
global) <- ByteString -> IO (ByteString, ByteString, Bool)
parse (String -> ByteString
B.pack String
ptn')
  Regex
regex <- ByteString -> IO Regex
compile ByteString
pat
  ByteString
res <-
    if Bool
global
      then Regex -> ByteString -> ByteString -> IO ByteString
replaceAll Regex
regex ByteString
rep (String -> ByteString
B.pack String
tgt')
      else Regex -> ByteString -> ByteString -> IO ByteString
replaceFirst Regex
regex ByteString
rep (String -> ByteString
B.pack String
tgt')
  Term -> IO Term
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expression -> Term
TeExpression (Bytes -> Expression
DataString (String -> Bytes
strToBts (ByteString -> String
B.unpack ByteString
res))))
  where
    parse :: B.ByteString -> IO (B.ByteString, B.ByteString, Bool)
    parse :: ByteString -> IO (ByteString, ByteString, Bool)
parse ByteString
input =
      case ByteString -> ByteString -> Maybe ByteString
B.stripPrefix ByteString
"s/" ByteString
input of
        Just ByteString
rest ->
          let parts :: [ByteString]
parts = Char -> ByteString -> [ByteString]
B.split Char
'/' ByteString
rest
           in case [ByteString]
parts of
                [ByteString
pat, ByteString
rep, ByteString
"g"] -> (ByteString, ByteString, Bool) -> IO (ByteString, ByteString, Bool)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString
pat, ByteString
rep, Bool
True)
                [ByteString
pat, ByteString
rep, ByteString
""] -> (ByteString, ByteString, Bool) -> IO (ByteString, ByteString, Bool)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString
pat, ByteString
rep, Bool
False)
                [ByteString
pat, ByteString
rep] -> (ByteString, ByteString, Bool) -> IO (ByteString, ByteString, Bool)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString
pat, ByteString
rep, Bool
False)
                [ByteString]
_ -> IOError -> IO (ByteString, ByteString, Bool)
forall e a. Exception e => e -> IO a
throwIO (String -> IOError
userError String
"sed pattern must be in format s/pat/rep/[g]")
        Maybe ByteString
_ -> IOError -> IO (ByteString, ByteString, Bool)
forall e a. Exception e => e -> IO a
throwIO (String -> IOError
userError String
"sed pattern must start with s/")
buildTermFromFunction String
"sed" [ExtraArgument]
_ Subst
_ Program
_ = IOError -> IO Term
forall e a. Exception e => e -> IO a
throwIO (String -> IOError
userError String
"Function sed() requires exactly 2 dataizable arguments")
buildTermFromFunction String
"random-string" [ExtraArgument
arg] Subst
subst Program
prog = do
  String
pat <- ExtraArgument -> Subst -> Program -> IO String
argToStrBytes ExtraArgument
arg Subst
subst Program
prog
  Set String
set <- IORef (Set String) -> IO (Set String)
forall a. IORef a -> IO a
readIORef IORef (Set String)
randomStrings
  String
str <- String -> Set String -> IO String
regenerate String
pat Set String
set
  Term -> IO Term
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expression -> Term
TeExpression (Bytes -> Expression
DataString (String -> Bytes
strToBts String
str)))
  where
    regenerate :: String -> Set String -> IO String
    regenerate :: String -> Set String -> IO String
regenerate String
pat Set String
set = do
      String
next <- String -> IO String
randomString String
pat
      if String
next String -> Set String -> Bool
forall a. Ord a => a -> Set a -> Bool
`Data.Set.member` Set String
set
        then String -> Set String -> IO String
regenerate String
pat Set String
set
        else do
          IORef (Set String) -> (Set String -> Set String) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef (Set String)
randomStrings (String -> Set String -> Set String
forall a. Ord a => a -> Set a -> Set a
Data.Set.insert String
next)
          String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure String
next
    randomString :: String -> IO String
    randomString :: String -> IO String
randomString [] = String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    randomString (Char
'%' : Char
ch : String
rest) = do
      String
rep <- case Char
ch of
        Char
'x' -> Int -> IO Char -> IO String
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
8 (IO Char -> IO String) -> IO Char -> IO String
forall a b. (a -> b) -> a -> b
$ do
          Int
v <- (Int, Int) -> IO Int
forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
randomRIO (Int
0, Int
15)
          Char -> IO Char
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Char
intToDigit Int
v)
        Char
'd' -> Int -> String
forall a. Show a => a -> String
show (Int -> String) -> IO Int -> IO String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, Int) -> IO Int
forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
randomRIO (Int
0 :: Int, Int
9999)
        Char
_ -> String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Char
'%', Char
ch]
      String
next <- String -> IO String
randomString String
rest
      String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String
rep String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
next)
    randomString (Char
ch : String
rest) = do
      String
rest' <- String -> IO String
randomString String
rest
      String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Char
ch Char -> String -> String
forall a. a -> [a] -> [a]
: String
rest')
buildTermFromFunction String
"random-string" [ExtraArgument]
_ Subst
_ Program
_ = IOError -> IO Term
forall e a. Exception e => e -> IO a
throwIO (String -> IOError
userError String
"Function random-string() requires exactly 1 dataizable argument")
buildTermFromFunction String
func [ExtraArgument]
_ Subst
_ Program
_ = IOError -> IO Term
forall e a. Exception e => e -> IO a
throwIO (String -> IOError
userError (String -> String -> String
forall r. PrintfType r => String -> r
printf String
"Function %s() is not supported or does not exist" String
func))