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

module Regexp where

import Control.Exception
import Data.Array (bounds, (!))
import qualified Data.ByteString.Char8 as B
import Data.Char (isDigit, ord)
import Data.Maybe (fromMaybe)
import Text.Read (readMaybe)
import qualified Text.Regex.PCRE.ByteString as R
import qualified Text.Regex.PCRE.Wrap as W

data Regex = Regex R.Regex R.Regex

compile :: B.ByteString -> IO Regex
compile :: ByteString -> IO Regex
compile ByteString
pat = do
  Regex
first <- ExecOption -> IO Regex
compileWith ExecOption
R.execBlank
  Regex
continuation <- ExecOption -> IO Regex
compileWith ExecOption
R.execNotBOL
  Regex -> IO Regex
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Regex -> Regex -> Regex
Regex Regex
first Regex
continuation)
  where
    compileWith :: ExecOption -> IO Regex
compileWith ExecOption
option = do
      Either (Int, [Char]) Regex
compiled <- CompOption
-> ExecOption -> ByteString -> IO (Either (Int, [Char]) Regex)
R.compile CompOption
R.compUTF8 ExecOption
option ([Char] -> ByteString
B.pack [Char]
"(*UCP)" ByteString -> ByteString -> ByteString
`B.append` ByteString
pat)
      case Either (Int, [Char]) Regex
compiled of
        Left (Int
_, [Char]
err) -> IOError -> IO Regex
forall e a. Exception e => e -> IO a
throwIO ([Char] -> IOError
userError ([Char]
"Regex compilation failed: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
err))
        Right Regex
regex -> Regex -> IO Regex
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Regex
regex

match :: B.ByteString -> B.ByteString -> IO Bool
match :: ByteString -> ByteString -> IO Bool
match ByteString
pattern ByteString
input = do
  Regex Regex
regex Regex
_ <- ByteString -> IO Regex
compile ByteString
pattern
  Either WrapError (Maybe (Array Int (Int, Int)))
res <- Regex
-> ByteString
-> IO (Either WrapError (Maybe (Array Int (Int, Int))))
R.execute Regex
regex ByteString
input
  case Either WrapError (Maybe (Array Int (Int, Int)))
res of
    Left WrapError
err -> WrapError -> IO Bool
forall err value. Show err => err -> IO value
regexExecutionFailed WrapError
err
    Right Maybe (Array Int (Int, Int))
Nothing -> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
    Right (Just Array Int (Int, Int)
_) -> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True

extractGroups :: Regex -> B.ByteString -> IO [B.ByteString]
extractGroups :: Regex -> ByteString -> IO [ByteString]
extractGroups (Regex Regex
regex Regex
_) ByteString
input = Regex -> ByteString -> IO [ByteString]
extractGroupsWith Regex
regex ByteString
input

extractGroupsWith :: R.Regex -> B.ByteString -> IO [B.ByteString]
extractGroupsWith :: Regex -> ByteString -> IO [ByteString]
extractGroupsWith Regex
regex ByteString
input = do
  Either WrapError (Maybe (Array Int (Int, Int)))
result <- Regex
-> ByteString
-> IO (Either WrapError (Maybe (Array Int (Int, Int))))
R.execute Regex
regex ByteString
input
  case Either WrapError (Maybe (Array Int (Int, Int)))
result of
    Left WrapError
err -> WrapError -> IO [ByteString]
forall err value. Show err => err -> IO value
regexExecutionFailed WrapError
err
    Right Maybe (Array Int (Int, Int))
Nothing -> [ByteString] -> IO [ByteString]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    Right (Just Array Int (Int, Int)
arr) ->
      let (Int
start, Int
end) = Array Int (Int, Int) -> (Int, Int)
forall i e. Array i e -> (i, i)
bounds Array Int (Int, Int)
arr
          groups :: [ByteString]
groups =
            [ let (Int
off, Int
len) = Array Int (Int, Int)
arr Array Int (Int, Int) -> Int -> (Int, Int)
forall i e. Ix i => Array i e -> i -> e
! Int
i
               in if Int
off Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== -Int
1 then ByteString
B.empty else Int -> ByteString -> ByteString
B.take Int
len (Int -> ByteString -> ByteString
B.drop Int
off ByteString
input)
            | Int
i <- [Int
start .. Int
end]
            ]
       in [ByteString] -> IO [ByteString]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [ByteString]
groups

substituteGroups :: B.ByteString -> [B.ByteString] -> B.ByteString
substituteGroups :: ByteString -> [ByteString] -> ByteString
substituteGroups ByteString
rep [ByteString]
groups = [ByteString] -> ByteString
B.concat ([Char] -> [ByteString]
go (ByteString -> [Char]
B.unpack ByteString
rep))
  where
    go :: String -> [B.ByteString]
    go :: [Char] -> [ByteString]
go [] = []
    go (Char
'$' : [Char]
rest) =
      let ([Char]
digits, [Char]
afterDigits) = (Char -> Bool) -> [Char] -> ([Char], [Char])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span Char -> Bool
isDigit [Char]
rest
       in if [Char] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Char]
digits
            then Char -> ByteString
B.singleton Char
'$' ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [Char] -> [ByteString]
go [Char]
rest
            else
              let val :: ByteString
val = case [Char] -> Maybe Integer
forall a. Read a => [Char] -> Maybe a
readMaybe [Char]
digits of
                    Maybe Integer
Nothing -> [Char] -> ByteString
B.pack (Char
'$' Char -> [Char] -> [Char]
forall a. a -> [a] -> [a]
: [Char]
digits)
                    Just Integer
idx -> ByteString -> Maybe ByteString -> ByteString
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> ByteString
B.pack (Char
'$' Char -> [Char] -> [Char]
forall a. a -> [a] -> [a]
: [Char]
digits)) (Integer -> [ByteString] -> Maybe ByteString
safeIndex Integer
idx [ByteString]
groups)
               in ByteString
val ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [Char] -> [ByteString]
go [Char]
afterDigits
    go (Char
c : [Char]
rest) = Char -> ByteString
B.singleton Char
c ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [Char] -> [ByteString]
go [Char]
rest
    safeIndex :: Integer -> [B.ByteString] -> Maybe B.ByteString
    safeIndex :: Integer -> [ByteString] -> Maybe ByteString
safeIndex Integer
i [ByteString]
xs
      | Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
0 Bool -> Bool -> Bool
&& Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([ByteString] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
xs) = ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ([ByteString]
xs [ByteString] -> Int -> ByteString
forall a. HasCallStack => [a] -> Int -> a
!! Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
i)
      | Bool
otherwise = Maybe ByteString
forall a. Maybe a
Nothing

replaceFirst :: Regex -> B.ByteString -> B.ByteString -> IO B.ByteString
replaceFirst :: Regex -> ByteString -> ByteString -> IO ByteString
replaceFirst (Regex Regex
regex Regex
_) ByteString
rep ByteString
input = do
  Either WrapError (Maybe (Array Int (Int, Int)))
result <- Regex
-> ByteString
-> IO (Either WrapError (Maybe (Array Int (Int, Int))))
R.execute Regex
regex ByteString
input
  case Either WrapError (Maybe (Array Int (Int, Int)))
result of
    Left WrapError
err -> WrapError -> IO ByteString
forall err value. Show err => err -> IO value
regexExecutionFailed WrapError
err
    Right Maybe (Array Int (Int, Int))
Nothing -> ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
input
    Right (Just Array Int (Int, Int)
arr) -> do
      [ByteString]
groups <- Regex -> ByteString -> IO [ByteString]
extractGroupsWith Regex
regex ByteString
input
      let (Int
off, Int
len) = Array Int (Int, Int)
arr Array Int (Int, Int) -> Int -> (Int, Int)
forall i e. Ix i => Array i e -> i -> e
! Int
0
          (ByteString
before, ByteString
rest) = Int -> ByteString -> (ByteString, ByteString)
B.splitAt Int
off ByteString
input
          (ByteString
_, ByteString
after) = Int -> ByteString -> (ByteString, ByteString)
B.splitAt Int
len ByteString
rest
          replacement :: ByteString
replacement = ByteString -> [ByteString] -> ByteString
substituteGroups ByteString
rep [ByteString]
groups
      ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> ByteString
B.concat [ByteString
before, ByteString
replacement, ByteString
after]

replaceAll :: Regex -> B.ByteString -> B.ByteString -> IO B.ByteString
replaceAll :: Regex -> ByteString -> ByteString -> IO ByteString
replaceAll (Regex Regex
first Regex
_) ByteString
rep ByteString
input = Int -> ByteString -> IO ByteString
go Int
0 ByteString
B.empty
  where
    go :: Int -> B.ByteString -> IO B.ByteString
    go :: Int -> ByteString -> IO ByteString
go Int
cursor ByteString
acc = do
      Either WrapError (Maybe [(Int, Int)])
result <- ByteString
-> (CStringLen -> IO (Either WrapError (Maybe [(Int, Int)])))
-> IO (Either WrapError (Maybe [(Int, Int)]))
forall a. ByteString -> (CStringLen -> IO a) -> IO a
B.useAsCStringLen ByteString
input (Int
-> Regex
-> CStringLen
-> IO (Either WrapError (Maybe [(Int, Int)]))
W.wrapMatch Int
cursor Regex
first)
      case Either WrapError (Maybe [(Int, Int)])
result of
        Left WrapError
err -> WrapError -> IO ByteString
forall err value. Show err => err -> IO value
regexExecutionFailed WrapError
err
        Right (Just offsets :: [(Int, Int)]
offsets@((Int
off, Int
end) : [(Int, Int)]
_)) -> do
          let len :: Int
len = Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
off
              before :: ByteString
before = Int -> ByteString -> ByteString
B.take (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
cursor) (Int -> ByteString -> ByteString
B.drop Int
cursor ByteString
input)
              rest :: ByteString
rest = Int -> ByteString -> ByteString
B.drop Int
end ByteString
input
              groups :: [ByteString]
groups = [if Int
start Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 then ByteString
B.empty else Int -> ByteString -> ByteString
B.take (Int
finish Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start) (Int -> ByteString -> ByteString
B.drop Int
start ByteString
input) | (Int
start, Int
finish) <- [(Int, Int)]
offsets]
              replacement :: ByteString
replacement = ByteString -> [ByteString] -> ByteString
substituteGroups ByteString
rep [ByteString]
groups
          if Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
            then
              let next :: ByteString
next = Int -> ByteString -> ByteString
B.take (ByteString -> Int
utf8Width ByteString
rest) ByteString
rest
               in if ByteString -> Bool
B.null ByteString
next
                    then ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> ByteString
B.concat [ByteString
acc, ByteString
before, ByteString
replacement]
                    else Int -> ByteString -> IO ByteString
go (Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ByteString -> Int
B.length ByteString
next) ([ByteString] -> ByteString
B.concat [ByteString
acc, ByteString
before, ByteString
replacement, ByteString
next])
            else Int -> ByteString -> IO ByteString
go Int
end ([ByteString] -> ByteString
B.concat [ByteString
acc, ByteString
before, ByteString
replacement])
        Right Maybe [(Int, Int)]
_ -> ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString -> ByteString
B.append ByteString
acc (Int -> ByteString -> ByteString
B.drop Int
cursor ByteString
input)
    utf8Width :: B.ByteString -> Int
    utf8Width :: ByteString -> Int
utf8Width ByteString
bytes =
      case ByteString -> Maybe (Char, ByteString)
B.uncons ByteString
bytes of
        Maybe (Char, ByteString)
Nothing -> Int
0
        Just (Char
lead, ByteString
rest)
          | Char -> Int
ord Char
lead Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0xC0 -> Int
1
          | Char -> Int
ord Char
lead Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0xE0 -> Int -> ByteString -> Int
continuationWidth Int
2 ByteString
rest
          | Char -> Int
ord Char
lead Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0xF0 -> Int -> ByteString -> Int
continuationWidth Int
3 ByteString
rest
          | Char -> Int
ord Char
lead Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0xF8 -> Int -> ByteString -> Int
continuationWidth Int
4 ByteString
rest
          | Bool
otherwise -> Int
1
    continuationWidth :: Int -> B.ByteString -> Int
    continuationWidth :: Int -> ByteString -> Int
continuationWidth Int
width ByteString
rest =
      if ByteString -> Int
B.length ByteString
rest Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Bool -> Bool -> Bool
&& (Char -> Bool) -> ByteString -> Bool
B.all Char -> Bool
isContinuation (Int -> ByteString -> ByteString
B.take (Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ByteString
rest)
        then Int
width
        else Int
1
    isContinuation :: Char -> Bool
    isContinuation :: Char -> Bool
isContinuation Char
byte = Char -> Int
ord Char
byte Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0x80 Bool -> Bool -> Bool
&& Char -> Int
ord Char
byte Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0xC0

regexExecutionFailed :: Show err => err -> IO value
regexExecutionFailed :: forall err value. Show err => err -> IO value
regexExecutionFailed err
err = IOError -> IO value
forall e a. Exception e => e -> IO a
throwIO ([Char] -> IOError
userError ([Char]
"Regex execution failed: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ err -> [Char]
forall a. Show a => a -> [Char]
show err
err))