{-# 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, when)
import qualified Data.ByteString.Char8 as B
import Data.Char (intToDigit)
import Data.Set (Set)
import qualified Data.Set
import GHC.IO (unsafePerformIO)
import Matcher
import Misc
import Numeric (showHex)
import Parser (parseAttributeThrows)
import Pretty
import Random (randomString)
import Regexp
import System.Random (randomRIO)
import Term
import Text.Printf (printf)
import qualified Yaml as Y

argToStrBytes :: Y.ExtraArgument -> Subst -> Program -> IO String
argToStrBytes :: ExtraArgument -> Subst -> Program -> IO String
argToStrBytes (Y.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 (Y.ArgExpression Expression
expr) Subst
subst Program
prog = do
  (TeBytes Bytes
bts) <- String -> [ExtraArgument] -> Subst -> Program -> IO Term
buildTermFromFunction String
"dataize" [Expression -> ExtraArgument
Y.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 -> [Y.ExtraArgument] -> Subst -> Program -> IO Term
buildTermFromFunction :: String -> [ExtraArgument] -> Subst -> Program -> IO Term
buildTermFromFunction String
"contextualize" [Y.ArgExpression Expression
expr, Y.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" [Y.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
  String
tau <- [String] -> IO String
randomTau [String]
attrs
  Term -> IO Term
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Attribute -> Term
TeAttribute (String -> Attribute
AtLabel String
tau))
  where
    argsToAttrs :: [Y.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
      Y.ArgExpression Expression
_ -> [ExtraArgument] -> IO [String]
argsToAttrs [ExtraArgument]
rest
      Y.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')
      Y.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')
      Y.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 :: [String] -> IO String
    randomTau :: [String] -> IO String
randomTau [String]
attrs = do
      String
tau <- String -> IO String
randomString String
"a🌵%d"
      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 [String] -> IO String
randomTau [String]
attrs else String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure String
tau
buildTermFromFunction String
"dataize" [Y.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" [Y.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]
args Subst
subst Program
prog = do
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([ExtraArgument] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExtraArgument]
args Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2) (IOError -> IO ()
forall e a. Exception e => e -> IO a
throwIO (String -> IOError
userError String
"Function sed() requires at least two arguments"))
  [ByteString]
args' <-
    (ExtraArgument -> IO ByteString)
-> [ExtraArgument] -> IO [ByteString]
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 -> do
          String
bts <- ExtraArgument -> Subst -> Program -> IO String
argToStrBytes ExtraArgument
arg Subst
subst Program
prog
          ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> ByteString
B.pack String
bts)
      )
      [ExtraArgument]
args
  ByteString
res <- ByteString -> [ByteString] -> IO ByteString
sed ([ByteString] -> ByteString
forall a. HasCallStack => [a] -> a
head [ByteString]
args') ([ByteString] -> [ByteString]
forall a. HasCallStack => [a] -> [a]
tail [ByteString]
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 (ByteString -> String
B.unpack ByteString
res))))
  where
    sed :: B.ByteString -> [B.ByteString] -> IO B.ByteString
    sed :: ByteString -> [ByteString] -> IO ByteString
sed ByteString
tgt [] = ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
tgt
    sed ByteString
tgt (ByteString
ptn : [ByteString]
ptns) = do
      (ByteString
pat, ByteString
rep, Bool
global) <- ByteString -> IO (ByteString, ByteString, Bool)
parse ByteString
ptn
      Regex
regex <- ByteString -> IO Regex
compile ByteString
pat
      ByteString
next <-
        if Bool
global
          then Regex -> ByteString -> ByteString -> IO ByteString
replaceAll Regex
regex ByteString
rep ByteString
tgt
          else Regex -> ByteString -> ByteString -> IO ByteString
replaceFirst Regex
regex ByteString
rep ByteString
tgt
      ByteString -> [ByteString] -> IO ByteString
sed ByteString
next [ByteString]
ptns
    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
body ->
          let (ByteString
pat, ByteString
rest) = ByteString -> ByteString -> Bool -> (ByteString, ByteString)
nextUntilSlash ByteString
body ByteString
B.empty Bool
False
              (ByteString
rep, ByteString
flag) = ByteString -> ByteString -> Bool -> (ByteString, ByteString)
nextUntilSlash ByteString
rest ByteString
B.empty Bool
True
           in case [ByteString
pat, ByteString
rep, ByteString
flag] 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]
_ -> 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/")
    -- Cut part from given string until regular slash.
    nextUntilSlash :: B.ByteString -> B.ByteString -> Bool -> (B.ByteString, B.ByteString)
    nextUntilSlash :: ByteString -> ByteString -> Bool -> (ByteString, ByteString)
nextUntilSlash ByteString
input ByteString
acc Bool
escape = case ByteString -> Maybe (Char, ByteString)
B.uncons ByteString
input of
      Maybe (Char, ByteString)
Nothing -> (ByteString
acc, ByteString
B.empty)
      Just (Char
h, ByteString
rest)
        | Char
h Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'\\' -> case ByteString -> Maybe (Char, ByteString)
B.uncons ByteString
rest of
            Just (Char
h', ByteString
rest') -> ByteString -> ByteString -> Bool -> (ByteString, ByteString)
nextUntilSlash ByteString
rest' ((if Bool
escape then ByteString
acc else ByteString -> Char -> ByteString
B.snoc ByteString
acc Char
'\\') ByteString -> ByteString -> ByteString
`B.append` Char -> ByteString
B.singleton Char
h') Bool
escape
            Maybe (Char, ByteString)
Nothing -> (if Bool
escape then ByteString
acc else ByteString -> Char -> ByteString
B.snoc ByteString
acc Char
'\\', ByteString
B.empty)
        | Char
h Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'/' -> (ByteString
acc, ByteString
rest)
        | Bool
otherwise -> ByteString -> ByteString -> Bool -> (ByteString, ByteString)
nextUntilSlash ByteString
rest (ByteString -> Char -> ByteString
B.snoc ByteString
acc Char
h) Bool
escape
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
  String
str <- String -> IO String
randomString String
pat
  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)))
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
"size" [Y.ArgBinding (BiMeta String
meta)] Subst
subst Program
_ = do
  [Binding]
bds <- Binding -> Subst -> IO [Binding]
buildBindingThrows (String -> Binding
BiMeta String
meta) Subst
subst
  Term -> IO Term
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expression -> Term
TeExpression (Bytes -> Expression
DataNumber (Double -> Bytes
numToBts (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Binding] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Binding]
bds)))))
buildTermFromFunction String
"size" [ExtraArgument]
_ Subst
_ Program
_ = IOError -> IO Term
forall e a. Exception e => e -> IO a
throwIO (String -> IOError
userError String
"Function size() requires exactly 1 meta binding")
buildTermFromFunction String
"tau" [Y.ArgExpression Expression
expr] Subst
subst Program
prog = do
  TeBytes Bytes
bts <- String -> [ExtraArgument] -> Subst -> Program -> IO Term
buildTermFromFunction String
"dataize" [Expression -> ExtraArgument
Y.ArgExpression Expression
expr] Subst
subst Program
prog
  Attribute
attr <- String -> IO Attribute
parseAttributeThrows (Bytes -> String
btsToUnescapedStr Bytes
bts)
  Term -> IO Term
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Attribute -> Term
TeAttribute Attribute
attr)
buildTermFromFunction String
"tau" [ExtraArgument]
_ Subst
_ Program
_ = IOError -> IO Term
forall e a. Exception e => e -> IO a
throwIO (String -> IOError
userError String
"Function tau() requires exactly 1 argument as expression")
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))