-- Copyright (c) 2023  Herbert Valerio Riedel <hvr@gnu.org>
--
--  This file is free software: you may copy, redistribute and/or modify it
--  under the terms of the GNU General Public License as published by the
--  Free Software Foundation, either version 2 of the License, or (at your
--  option) any later version.
--
--  This file is distributed in the hope that it will be useful, but
--  WITHOUT ANY WARRANTY; without even the implied warranty of
--  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
--  General Public License for more details.
--
--  You should have received a copy of the GNU General Public License
--  along with this program (see `LICENSE`).  If not, see
--  <https://www.gnu.org/licenses/old-licenses/gpl-2.0.html>.

{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE Trustworthy                #-}

-- |
-- Copyright: © Herbert Valerio Riedel 2023
-- SPDX-License-Identifier: GPL-2.0-or-later
--
-- LDAP OID Helpers
--
-- This module provides helpers for dealing with the representation of /Object Identifiers/ (OID) in LDAP.
--
-- @since 0.1.2
module LDAPv3.OID
   ( -- * Textually encoded OIDs
     LDAPOID
   , OID(OID)

     -- * Binary encoded OIDs
   , OBJECT_IDENTIFIER
   , object_identifier'toOID
   , object_identifier'fromOID

   , object_identifier'toBin
   , object_identifier'fromBin

     -- * Convenience helpers
   , IsWellFormedOid(isWellFormedOid)
   ) where

import           Common                  hiding (Option, many, option, some, (<|>))
import           Data.ASN1
import           Data.ASN1.Prim
import           LDAPv3.StringRepr.Class

import qualified Data.Binary             as Bin
import qualified Data.ByteString.Builder as BSB
import qualified Data.ByteString.Char8   as BSC
import qualified Data.ByteString.Lazy    as BSL
import qualified Data.ByteString.Short   as SBS
import           Data.List
import qualified Data.Text               as T
import qualified Data.Text.Lazy          as TL
import           Data.Text.Lazy.Builder  as B
import qualified Data.Text.Short         as TS
import           Numeric                 (showHex)
import           Text.Parsec             as P

-- | Typeclass for 'isWellFormedOid' operation
--
-- @since 0.1.2
class IsWellFormedOid t where
  -- | Determine whether OID representation is deemed well-formed
  --
  -- An OID is considered well-formed /iff/ it has
  --
  --  * at least two arcs,
  --  * the first arc is one of @0@, @1@, or @2@, and
  --  * if the first arc is /not/ @2@, the second arc value is within the range @[0 .. 39]@.
  --
  -- Additionally, for string types the IETF-style ASCII dot notation with normalized (i.e. without redundant leading
  -- zeros) decimal numbers is expected (e.g. @1.23.456.7.890@) as expressed by the @numericoid@ ABNF production shown
  -- below:
  --
  -- > numericoid = number 1*( DOT number )
  -- > number  = DIGIT / ( LDIGIT 1*DIGIT )
  -- > DIGIT   = %x30 / LDIGIT       ; "0"-"9"
  -- > LDIGIT  = %x31-39             ; "1"-"9"
  --
  -- @since 0.1.2
  isWellFormedOid :: t -> Bool


{- | Object identifier  (<https://tools.ietf.org/html/rfc4511#section-4.1.2 RFC4511 Section 4.1.2>)

> LDAPOID ::= OCTET STRING -- Constrained to <numericoid>
>                          -- [RFC4512]

@since 0.1.0
-}
type LDAPOID = OID


{- | Numeric Object Identifier (OID)

> numericoid = number 1*( DOT number )
> number  = DIGIT / ( LDIGIT 1*DIGIT )
> DIGIT   = %x30 / LDIGIT       ; "0"-"9"
> LDIGIT  = %x31-39             ; "1"-"9"

NB: The current type definition and its 'StringRepr' instance currently allows to represent and parse more than the ABNF
    described above; moreover, the ABNF is also more liberal as it doesn't express the constraints imposed upon the
    first two arcs by @X.660@ and @ASN.1@. See also 'isWellFormedOid'.

@since 0.1.0
-}
newtype OID = OID (NonEmpty Natural)
  deriving (Eq,Ord,Show,NFData)

instance Newtype OID (NonEmpty Natural)

instance ASN1 OID where
  asn1defTag _   = asn1defTag (Proxy :: Proxy OCTET_STRING)
  asn1encode oid = asn1encode (BSC.pack (s'OID oid))
  asn1decode     = asn1decodeParsec "OID" p'OID

instance StringRepr OID where
  asBuilder       = B.fromString . s'OID
  renderShortText = TS.fromString . s'OID
  asParsec        = p'OID

-- Ideally this instance will become redundant/trivial in the future as it currently papers over an inadequacy in the
-- current definition of the 'OID' type
instance IsWellFormedOid OID where
  isWellFormedOid (OID s) = case s of
    0 :| (y:_) | y < 40 -> True
    1 :| (y:_) | y < 40 -> True
    2 :| (_:_)          -> True
    _                   -> False

-- the main cost is dealing with formatting the 'Natural' components and it's not obvious if it's worth the complexity
-- optimizing a more direct path for 'ShortText'
s'OID :: OID -> String
s'OID (OID (x:|xs)) = intercalate "." (map show (x:xs))

p'OID :: Stream s Identity Char => Parsec s () OID
p'OID = p'numericoid
  where
    -- numericoid = number 1*( DOT number )
    p'numericoid = OID <$> (p'number `sepBy1'` char '.')

    -- number  = DIGIT / ( LDIGIT 1*DIGIT )
    -- DIGIT   = %x30 / LDIGIT       ; "0"-"9"
    -- LDIGIT  = %x31-39             ; "1"-"9"
    p'number = do
      ldigit <- digit
      if ldigit == '0'
         then pure 0
         else read . (ldigit:) <$> many digit

{-# INLINE wf'OID #-}
wf'OID :: String -> Bool
wf'OID = \case
    '0':'.':rest -> go01 rest
    '1':'.':rest -> go01 rest
    '2':'.':rest -> go rest
    _            -> False
  where
    -- enforce 2nd arc to be within (canonically represented) [0..39] range
    go01 [d1]    | isD d1                       = True
    go01 [d1,d2] | d1 >= '1', d1 <= '3', isD d2 = True
    go01 (d1:'.':rest)    | isD d1                       = go rest
    go01 (d1:d2:'.':rest) | d1 >= '1', d1 <= '3', isD d2 = go rest
    go01 _ = False

    -- subsequent arcs without upper bounds
    go (c:rest)
      | c == '0' = case rest of
                    []        -> True
                    '.':rest' -> go rest'
                    _         -> False
      | isNZD c = case dropWhile isD rest of
                    []        -> True
                    '.':rest' -> go rest'
                    _:_       -> False
    go _ = False

    isNZD c = c >= '1' && c <= '9'
    isD c   = c >= '0' && c <= '9'

-- the instances below work best if the compiler manages to optimize away the intermediate @[Char]@ ...

instance IsWellFormedOid ShortText where
    isWellFormedOid = wf'OID . TS.unpack

instance IsWellFormedOid T.Text where
    isWellFormedOid = wf'OID . T.unpack

instance IsWellFormedOid TL.Text where
    isWellFormedOid = wf'OID . TL.unpack


-- | ASN.1 @OBJECT IDENTIFIER@
--
-- The 'OID' type uses the textual LDAP encoding when converted to/from ASN.1 whereas this type provides the proper ASN.1 encoding as defined per X.690 section 8.19 (accessible via its 'Binary' instance).
--
-- @since 0.1.2
newtype OBJECT_IDENTIFIER = OID_ SBS.ShortByteString {- content encoded as per X.690 sec. 8.19 -}
  deriving (Eq,NFData)

instance Show OBJECT_IDENTIFIER where
    showsPrec _ (OID_ z) = ("OID<"++) . (\s -> foldr hex8 s (SBS.unpack z)) . ('>':)
      where
        hex8 :: Word8 -> ShowS
        hex8 x
          | x < 0x10  = ('0':) . showHex x
          | otherwise = showHex x

-- | Lexicographic ordering
instance Ord OBJECT_IDENTIFIER where
  compare = compareSubIds

instance ASN1 OBJECT_IDENTIFIER where
  asn1defTag _ = Universal 6
  asn1encode (OID_ sbs) = retag (Universal 6) $ asn1encode sbs
  asn1decode = implicit (Universal 6)
               (asn1decode `transformVia`
                (maybe (Left "not well-formed OBJECT IDENTIFIER") Right . object_identifier'fromBin))

instance StringRepr OBJECT_IDENTIFIER where
  asBuilder       = asBuilder . object_identifier'toOID
  renderShortText = renderShortText . object_identifier'toOID
  asParsec        = do
    x <- p'OID
    case object_identifier'fromOID x of
      Nothing -> fail "invalid top-level arcs"
      Just y  -> pure y

-- | Encodes as ASN.1 BER\/DER with @UNIVERSAL 6@ tag as per X.690 section 8.19
instance Bin.Binary OBJECT_IDENTIFIER where
  get = toBinaryGet asn1decode
  put = void . toBinaryPut . asn1encode

-- | Trivial instance as 'OBJECT_IDENTIFIER' values are always well-formed by construction
instance IsWellFormedOid OBJECT_IDENTIFIER where
  isWellFormedOid = const True

-- | Encode as raw ASN.1 BER\/DER value (i.e. without tag & length)
--
-- NB: As this function simply returns the internal representation this operation has zero cost.
--
-- @since 0.1.2
object_identifier'toBin :: OBJECT_IDENTIFIER -> SBS.ShortByteString
object_identifier'toBin (OID_ sbs) = sbs

-- | Decode from raw ASN.1 BER\/DER value (i.e. without tag & length)
--
-- All byte sequences are deemed well-formed raw ASN.1 OID encodings that satisfy the simple rules below (which ought to result in the same syntax as the rules specified in X.690 section 8.19.):
--
--  * The sequence must end with an octet with a value below @0x80@ (i.e. unset MSB), and
--  * any @0x80@ octet must be directly preceded by an octet which must have a value equal or greater than @0x80@ (i.e. set MSB).
--
-- In case these rules are not satisfied this function returns 'Nothing'.
--
-- NB: As this encoding matches the internal representation the resulting 'OBJECT_IDENTIFIER' merely @newtype@-wraps the input argument on success.
--
-- @since 0.1.2
object_identifier'fromBin :: SBS.ShortByteString -> Maybe OBJECT_IDENTIFIER
object_identifier'fromBin z
  | isValid   = Just (OID_ z)
  | otherwise = Nothing
  where
    isValid = case SBS.unpack z of
      [] -> False
      xs -> go 0x00 xs

    go pre []       = pre < 0x80
    go pre (0x80:xs)
      | pre >= 0x80 = go 0x80 xs
      | otherwise   = False
    go _ (x:xs)     = go x xs

-- | Convert 'OBJECT_IDENTIFIER' into 'OID' representation
--
-- @since 0.1.2
object_identifier'toOID :: OBJECT_IDENTIFIER -> OID
object_identifier'toOID oid = case decodeSubIds oid of
  [] -> error "the impossible just happened: internal invariant of OBJECT_IDENTIFIER broken"
  (i0:is)
    | i0 < 40   -> OID (0 :| (i0 : is))
    | i0 < 80   -> OID (1 :| (i0-40 : is))
    | otherwise -> OID (2 :| (i0-80 : is))

-- NB: the next major version shall avoid the 'Maybe' in the typesig
-- | Try to 'OID' representation into 'OBJECT_IDENTIFIER' representation
--
-- NB: This will return 'Nothing' /iff/ 'isWellFormedOid' returns 'False' on the input argument.
--
-- @since 0.1.2
object_identifier'fromOID :: OID -> Maybe OBJECT_IDENTIFIER
object_identifier'fromOID (OID s) = case s of
  0 :| (y:rest) | y < 40 -> Just $ encodeSubIds (y :| rest)
  1 :| (y:rest) | y < 40 -> Just $ encodeSubIds (y+40 :| rest)
  2 :| (y:rest)          -> Just $ encodeSubIds (y+80 :| rest)
  _                      -> Nothing -- invalid OID

encodeSubIds :: NonEmpty Natural -> OBJECT_IDENTIFIER
encodeSubIds (z:|zs) = OID_ $ SBS.toShort $ BSL.toStrict $ BSB.toLazyByteString $ mconcat (map subid (z:zs))
  where
    subid :: Natural -> BSB.Builder
    subid x
      | x < 0x100000000 {- i.e. 2^32 -} = encodeWord32 False (fromIntegral x) -- shortcut
      | otherwise = let (x',x'') = fmap fromIntegral (quotRem x 0x10000000) -- NB: 0x80^4, *not* 2^32
                    in subid1 x' `mappend` encodeWord32 True x''

    subid1 :: Natural -> BSB.Builder
    subid1 0 = mempty -- NB: never reached due to shortcut
    subid1 x
      | x < 0x80 = BSB.word8 (fromIntegral x .|. 0x80) -- shortcut
      | otherwise = let (x',x'') = fmap fromIntegral (quotRem x 0x80)
                    in subid1 x' `mappend` (BSB.word8 (x'' .|. 0x80))

    -- fast-path for 32bit values encoded in up to 5 octects;
    -- if enabled, pre-pad with 0x80 octects to 4 octets
    encodeWord32 :: Bool -> Word32 -> BSB.Builder
    encodeWord32 dopad w
      | w < 0x80     = pad' (BSB.word16BE 0x8080 <> BSB.word8 0x80) $
            BSB.word8 (fromIntegral w)
      | w < 0x4000   = pad' (BSB.word16BE 0x8080) $
            BSB.word8 (fromIntegral (w `unsafeShiftR`  7) .|. 0x80) <>
            BSB.word8 (fromIntegral w .&. 0x7f)
      | w < 0x200000 = pad' (BSB.word8 0x80) $
            BSB.word8 (fromIntegral (w `unsafeShiftR` 14) .|. 0x80) <>
            BSB.word8 (fromIntegral (w `unsafeShiftR`  7) .|. 0x80) <>
            BSB.word8 (fromIntegral w .&. 0x7f)
      | w < 0x10000000 =
            BSB.word8 (fromIntegral (w `unsafeShiftR` 21) .|. 0x80) <>
            BSB.word8 (fromIntegral (w `unsafeShiftR` 14) .|. 0x80) <>
            BSB.word8 (fromIntegral (w `unsafeShiftR`  7) .|. 0x80) <>
            BSB.word8 (fromIntegral w .&. 0x7f)
      | dopad = error "the impossible happened (encodeWord32)"
      | otherwise =
            BSB.word8 (fromIntegral (w `unsafeShiftR` 28) .|. 0x80) <>
            BSB.word8 (fromIntegral (w `unsafeShiftR` 21) .|. 0x80) <>
            BSB.word8 (fromIntegral (w `unsafeShiftR` 14) .|. 0x80) <>
            BSB.word8 (fromIntegral (w `unsafeShiftR`  7) .|. 0x80) <>
            BSB.word8 (fromIntegral w .&. 0x7f)
      where
        pad' thepad x
          | dopad     = thepad <> x
          | otherwise = x


decodeSubIds :: OBJECT_IDENTIFIER -> [Natural]
decodeSubIds (OID_ b) = go0 (SBS.unpack b)
  where
    go0 [] = []
    go0 (x:rest)
      | x < 0x80 = fromIntegral x : go0 rest
      | otherwise = go1 (shift7 $ stripMsb x) rest

    go1 _ [] = error "the impossible just happened: internal invariant of OBJECT_IDENTIFIER broken"
    go1 acc (x:rest)
      | x < 0x80  = (acc + fromIntegral x) : go0 rest
      | otherwise = go1 (shift7 $ acc + stripMsb x) rest

    stripMsb :: Word8 -> Natural
    stripMsb x = fromIntegral (x .&. 0x7f)

    shift7 :: Natural -> Natural
    shift7 x = x `unsafeShiftL` 7

-- efficient lexicographic streaming
--
-- This is supposed to be semantically equivalent to
--
-- > compareSubIds x y = compare (decodeSubIds x) (decodeSubIds y)
--
-- NB: this implementation relies on 'mappend' over 'Ordering' shortcutting
compareSubIds :: OBJECT_IDENTIFIER -> OBJECT_IDENTIFIER -> Ordering
compareSubIds x y | x == y = EQ -- fast shortcut
compareSubIds (OID_ bx) (OID_ by) = go EQ (SBS.unpack bx) (SBS.unpack by)
  where
    go c (x:xs) (y:ys)
      | finX, finY = c <> compare x y <> go EQ xs ys
      | finX       = LT
      | finY       = GT
      | otherwise  = go (c <> compare x y) xs ys
      where
        finX = x < 0x80
        finY = y < 0x80

    go EQ (_:_) [] = GT
    go EQ [] (_:_) = LT
    go EQ [] []    = EQ -- actually not reachable due to '==' short-cut

    go c xs ys = error $ "the impossible just happened: compareSubIds " ++ show (bx,by,c,xs,ys)