{-# 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.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 Yaml

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
  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 :: [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 :: [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" [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
  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" [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" [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
  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))