{-# LANGUAGE LambdaCase #-}
module KeyDB
    ( TrustMap
    , SigAndTrust
    , SubKey(..)
    , KeyData(..)
    , KeyDB(..)
    , KeyGrip(..)
    , emptyKeyDB
    , keyData
    , kkData
    , lookupKeyData
    , lookupByGrip
    , associatedKeys
    , fingerprintGrip
    , smallprGrip
    , transmute
    , transmuteAt
    , alterKeyDB
    , mergeKeyDB
    , mapKeyDB
    -- These probably don't belong here
    , selectKey0
    , flattenTop
    , flattenAllUids
    , flattenSub
    , sortByHint
    , flattenKeys
    , flattenFiltered
    , UidString(..)
    , buildGripMap
    ) where

import Control.Monad
import Data.Binary
import qualified Data.ByteString as S
import qualified Data.ByteString.Lazy as L
import Data.Functor
import Data.List
import qualified Data.Map.Strict as Map
import Data.Maybe
import Data.OpenPGP
import Data.Ord
import Foreign.Storable
import Text.Read

import Data.List.Merge
import Data.OpenPGP.Util
import qualified IntMapClass as I
         ;import IntMapClass (IMap)
import KeyRing.Types

type TrustMap = Map.Map FilePath Packet
type SigAndTrust = ( MappedPacket
                   , TrustMap     ) -- trust packets
data SubKey = SubKey MappedPacket [SigAndTrust] deriving Show

data UidString = UidString { unUidString :: String } deriving (Show, Eq, Ord)

-- | This is a GPG Identity which includes a master key and all its UIDs and
-- subkeys and associated signatures.
data KeyData = KeyData
    { keyMappedPacket :: MappedPacket                               -- main key
    , keySigAndTrusts :: [SigAndTrust]                              -- sigs on main key
    , keyUids         :: (Map.Map UidString ([SigAndTrust],OriginMap)) -- uids
    , keySubKeys      :: (Map.Map KeyKey SubKey)                    -- subkeys
    } deriving Show


newtype KeyGrip = KeyInt Int
 deriving Eq

fingerprintGrip :: Fingerprint -> KeyGrip
fingerprintGrip (Fingerprint bs) =
    -- case decode $ L.fromStrict $ S.drop (S.length bs - sizeOf (0::Int)) bs of
    -- -- The above was removed because Int is encoded as 8 bytes even when we are
    -- -- using 32-bit GHC.
    -- Presumably, the extra 4 bytes will be truncated.
    case S.length bs of
        -- v5 from the front
        32 -> case decode $ L.fromStrict bs of
            i -> KeyInt i
        -- v4 from the back
        l -> case decode $ L.fromStrict $ S.drop (l - 8) bs of
            i -> KeyInt i

smallprGrip :: String -> Maybe KeyGrip
smallprGrip pr = case length pr of
    64 -> KeyInt <$> readMaybe ("0x" ++ take (2 * sizeOf (0::Int)) pr)
    l  -> KeyInt <$> readMaybe ("0x" ++ drop (l - 2 * sizeOf (0::Int)) pr)

data KeyDB = KeyDB
    { byKeyKey :: Map.Map KeyKey KeyData
    , byGrip   :: IMap KeyGrip [KeyKey]
    } deriving Show


-- | TODO: This is an optimization to legacy (pre-KeyDB) code.  Ultimately it
-- should be unneccessary.
buildGripMap :: [Packet] -> IMap KeyGrip [Packet]
buildGripMap ps = foldr go I.empty ps
 where
    go pkt m =  I.alter (\case Just ks -> Just (pkt:ks)
                               Nothing -> Just [pkt])
                        (fingerprintGrip . fingerprint $ pkt)
                        m

emptyKeyDB :: KeyDB
emptyKeyDB = KeyDB { byKeyKey = Map.empty, byGrip = I.empty }

keyData :: KeyDB -> [KeyData]
keyData db = Map.elems (byKeyKey db)

kkData :: KeyDB -> [(KeyKey, KeyData)]
kkData db = Map.toList (byKeyKey db)

lookupKeyData :: KeyKey -> KeyDB -> Maybe KeyData
lookupKeyData kk db = Map.lookup kk (byKeyKey db)

lookupByGrip :: KeyGrip -> KeyDB -> [(MappedPacket,KeyData)]
lookupByGrip k db = do
    kk <- concat $ maybeToList $ I.lookup k (byGrip db)
    case Map.lookup kk (byKeyKey db) of
        Just kd | fingerprintGrip (fingerprint (packet $ keyMappedPacket kd)) == k -> [(keyMappedPacket kd, kd)]
                | otherwise -> do
                    sub <- associatedKeys kd
                    guard (mpGrip sub == k)
                    [ (sub, kd) ]
        Nothing -> do
            kd <- Map.elems (byKeyKey db)
            sub <- associatedKeys kd
            guard (mpGrip sub == k)
            [ (sub, kd) ]

transmute :: (Monad m, Monad kiki, Traversable kiki) =>
       ((KeyData, [info]) -> opcode -> m (kiki (KeyData, [info]))) -- ^ interpreter
       -> (KeyData -> [opcode])                                    -- ^ instructions
       -> KeyDB                                                    -- ^ initial state
       -> m (kiki (KeyDB, [info]))
transmute perform update db = do
    let performAll kd = foldM (\kkd op -> join <$> mapM (`perform` op) kkd)
                              (pure (kd,[]))
                              (update kd)
    r <- sequenceA <$> mapM performAll (byKeyKey db)
    return $ r <&> \bkk -> ( db { byKeyKey = fst <$> bkk }
                           -- Note: We currently leave deleted-keys in the byGrip map.
                           , concatMap snd $ Map.elems bkk )

mpGrip :: MappedPacket -> KeyGrip
mpGrip mp = fingerprintGrip $ fingerprint $ packet mp

associatedKeys :: KeyData -> [MappedPacket]
associatedKeys kd = keyMappedPacket kd : [ k | SubKey k _ <- Map.elems (keySubKeys kd) ]

alterKeyDB :: (Maybe KeyData -> Maybe KeyData) -> KeyKey -> KeyDB -> KeyDB
alterKeyDB update kk db = db
    { byKeyKey = Map.alter update kk (byKeyKey db)
    , byGrip = {- case Map.lookup kk (byKeyKey db) of
        Just _  -> byGrip db
        Nothing -> -} case update Nothing of
            Just kd -> let go g m = I.alter (\case Nothing  -> Just [kk]
                                                   Just kks -> Just $ mergeL [kk] kks)
                                            g
                                            m
                        in foldr go (byGrip db) $ map mpGrip $ filter (isKey . packet) $ associatedKeys kd
            Nothing -> byGrip db
    }

transmuteAt :: ( Monad m
               , Functor kiki
               ) => (Maybe KeyData -> m (kiki (KeyData,[info]))) -> KeyKey -> KeyDB -> m (kiki (KeyDB,[info]))
transmuteAt go kk db = do
    kdr <- go (Map.lookup kk $ byKeyKey db)
    return $ kdr <&> \(kd',rrs) -> ( alterKeyDB (const $ Just kd') kk db
                                   , rrs )

mergeKeyDB :: (KeyData -> KeyData -> KeyData) -> KeyDB -> KeyDB -> KeyDB
mergeKeyDB mergeKeyData db dbtrans =
    KeyDB { byKeyKey = Map.unionWith mergeKeyData (byKeyKey db) (byKeyKey dbtrans)
          , byGrip   = I.unionWith mergeL (byGrip db) (byGrip dbtrans)
          }

mapKeyDB :: Monad m => (KeyData -> m KeyData) -> KeyDB -> m KeyDB
mapKeyDB f db = fmap (\m -> db { byKeyKey = m }) $ mapM f (byKeyKey db)

selectKey0 :: Bool -> (KeySpec,Maybe String) -> KeyDB -> Maybe Packet
selectKey0 wantPublic (spec,mtag) db = do
    let Message ps = flattenKeys wantPublic $ byKeyKey db
        ys = snd $ seek_key spec ps
    flip (maybe (listToMaybe ys)) mtag $ \tag -> do
      case ys of
          y:ys1 -> listToMaybe $ snd $ seek_key (KeyTag y tag) ys1
          []    -> Nothing


flattenKeys :: Bool -> Map.Map KeyKey KeyData -> Message
flattenKeys isPublic db = Message $ concatMap (map packet . flattenTop "" isPublic . snd)
                                              (prefilter . Map.assocs $ db)
 where
    prefilter = if isPublic then id else filter isSecret
                 where
                  isSecret (_,(KeyData
                                  (MappedPacket { packet=(SecretKeyPacket {})})
                                  _
                                  _
                                  _)) = True
                  isSecret _          = False

flattenUid :: FilePath -> Bool -> (UidString,([SigAndTrust],OriginMap)) -> [MappedPacket]
flattenUid fname ispub (UidString str,(sigs,om)) =
    (mappedPacket "" $ UserIDPacket str) {locations=om} : concatSort fname head (unsig fname ispub) sigs

flattenSub :: FilePath -> Bool -> SubKey -> [MappedPacket]
flattenSub fname ispub (SubKey key sigs) = unk ispub key: concatSort fname head (unsig fname ispub) sigs

flattenAllUids :: FilePath -> Bool -> Map.Map UidString ([SigAndTrust],OriginMap) -> [MappedPacket]
flattenAllUids fname ispub uids =
    concatSort fname head (flattenUid fname ispub) (Map.assocs uids)

flattenTop :: FilePath -> Bool -> KeyData -> [MappedPacket]
flattenTop fname ispub (KeyData key sigs uids subkeys) =
    unk ispub key :
            ( flattenAllUids fname ispub uids
             ++ concatSort fname head (flattenSub fname ispub) (Map.elems subkeys))

sortByHint :: FilePath -> (a -> MappedPacket) -> [a] -> [a]
sortByHint fname f = sortBy (comparing gethint)
  where
    gethint = maybe defnum originalNum . Map.lookup fname . locations . f
    defnum = -1

concatSort ::
  FilePath -> ([a] -> MappedPacket) -> (b -> [a]) -> [b] -> [a]
concatSort fname getp f = concat . sortByHint fname getp . map f

unk :: Bool -> MappedPacket -> MappedPacket
unk isPublic = if isPublic then toPacket secretToPublic else id
                where toPacket f mp@(MappedPacket {packet=p}) = mp {packet=(f p)}


unsig :: FilePath -> Bool -> SigAndTrust -> [MappedPacket]
unsig fname isPublic (sig,trustmap) =
    sig : map (asMapped (-1)) ( take 1 . Map.elems $ Map.filterWithKey f trustmap)
  where
    f n _ = n==fname -- && trace ("fname=n="++show n) True
    asMapped n p = let m = mappedPacket fname p
                   in m { locations = fmap (\x->x {originalNum=n}) (locations m) }

flattenFiltered :: Bool -> (KeyData -> Bool) -> KeyDB -> Message
flattenFiltered wantPublic pred db = flattenKeys wantPublic $ Map.filter pred (byKeyKey db)