-- TK.hs: OpenPGP (RFC9580) transferable key data type
-- Copyright © 2012-2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

module Codec.Encryption.OpenPGP.Types.Internal.TK where

import Control.Arrow ((&&&))
import Control.Comonad (Comonad (..))
import Control.Lens
    ( folded
    , makeLenses
    , to
    , view
    , (^.)
    , (^..)
    , _1
    )
import qualified Data.Aeson.TH as ATH
import Data.Bifunctor (first)
import qualified Data.ByteString.Lazy as BL
import Data.Data (Data)
import Data.Function (on)
import qualified Data.HashMap.Lazy as HashMap
import Data.IxSet.Typed (IxSet)
import qualified Data.IxSet.Typed as IxSet
import Data.Kind (Type)
import Data.List (find, nub, sort, sortOn)
import qualified Data.List.NonEmpty as NE
import qualified Data.Map as Map
import Data.Ord (comparing)
import Data.Semigroup (Semigroup (..))
import Data.Text (Text)
import Data.Typeable (Typeable)
import Data.Word (Word8)
import GHC.Generics (Generic)

import Codec.Encryption.OpenPGP.Types.Internal.Base
import Codec.Encryption.OpenPGP.Types.Internal.PKITypes
import Codec.Encryption.OpenPGP.Types.Internal.Pkt

-- | Zipper for navigating a list of packets with position context
data PacketZipper a
    = PacketZipper
    { _zpBefore :: [a] -- packets consumed before focus
    , _zpCurrent :: a -- current packet under focus
    , _zpAfter :: [a] -- packets remaining after focus
    }
    deriving (Data, Eq, Functor, Generic, Show, Typeable)

instance Comonad PacketZipper where
    -- Extract: read the packet currently under focus
    extract (PacketZipper _ cur _) = cur

    -- Duplicate: produce a zipper-of-zippers,
    -- each focused on a different position in the original stream.
    duplicate z@(PacketZipper before cur after) =
        PacketZipper
            [ PacketZipper
                (take i before)
                (before !! i)
                (drop (i + 1) before ++ [cur] ++ after)
            | i <- [0 .. length before - 1]
            ]
            z
            [ PacketZipper
                (before ++ [cur] ++ take i after)
                (after !! i)
                (drop (i + 1) after)
            | i <- [0 .. length after - 1]
            ]

-- | Create a zipper from a list, starting at the first element
zFromList :: [a] -> Maybe (PacketZipper a)
zFromList [] = Nothing
zFromList (x : xs) = Just (PacketZipper [] x xs)

-- | Extract position in packet sequence
zPosition :: PacketZipper a -> Int
zPosition (PacketZipper before _ _) = length before

-- | Move to the next packet in the sequence
zMoveNext :: PacketZipper a -> Maybe (PacketZipper a)
zMoveNext (PacketZipper before current []) = Nothing
zMoveNext (PacketZipper before current (x : xs)) =
    Just (PacketZipper (before ++ [current]) x xs)

-- | Get the remaining packets (current + after)
zRemaining :: PacketZipper a -> [a]
zRemaining (PacketZipper _ current after) = current : after

-- | Reconstruct the full list from a zipper
zToList :: PacketZipper a -> [a]
zToList (PacketZipper before current after) = before ++ [current] ++ after

data TKUnknown
    = TKUnknown
    { _tkuKey :: (SomePKPayload, Maybe SKAddendum)
    , _tkuRevs :: [SignaturePayload]
    , _tkuUIDs :: [(Text, [SignaturePayload])]
    , _tkuUAts :: [([UserAttrSubPacket], [SignaturePayload])]
    , _tkuSubs :: [(Pkt, [SignaturePayload])]
    }
    deriving (Data, Eq, Generic, Show, Typeable)

data TKKind
    = PublicTK
    | SecretTK
    deriving (Data, Eq, Generic, Ord, Show, Typeable)

type family TKKindToKeyPktKind (k :: TKKind) :: KeyPktKind where
    TKKindToKeyPktKind 'PublicTK = 'PublicPkt
    TKKindToKeyPktKind 'SecretTK = 'SecretPkt

data TK (k :: TKKind)
    = TK
    { _tkPrimaryKey :: KeyPkt (TKKindToKeyPktKind k)
    , _tkRevs :: [SignaturePayload]
    , _tkUIDs :: [(Text, [SignaturePayload])]
    , _tkUAts :: [([UserAttrSubPacket], [SignaturePayload])]
    , _tkSubs :: [(KeyPkt (TKKindToKeyPktKind k), [SignaturePayload])]
    }
    deriving (Eq, Show)

deriving instance
    (Data (KeyPkt (TKKindToKeyPktKind k)), Typeable k) => Data (TK k)

instance Ord (TK k) where
    compare = comparing _tkPrimaryKey

data SomeTK where
    SomePublicTK :: TK 'PublicTK -> SomeTK
    SomeSecretTK :: TK 'SecretTK -> SomeTK

deriving instance Show SomeTK

instance Eq SomeTK where
    left == right = someTKToUnknown left == someTKToUnknown right

data TKConversionError
    = PublicSubkeyHasPrimaryRole
    | SecretSubkeyHasPrimaryRole
    | ExpectedPublicSubkeyPacket Word8
    | ExpectedSecretSubkeyPacket Word8
    deriving (Eq, Show)

renderTKConversionError :: TKConversionError -> String
renderTKConversionError PublicSubkeyHasPrimaryRole = "public subkey has primary-key role"
renderTKConversionError SecretSubkeyHasPrimaryRole = "secret subkey has primary-key role"
renderTKConversionError (ExpectedPublicSubkeyPacket tagValue) =
    "expected public subkey, got packet tag " ++ show tagValue
renderTKConversionError (ExpectedSecretSubkeyPacket tagValue) =
    "expected secret subkey, got packet tag " ++ show tagValue

tkToUnknown :: TK k -> TKUnknown
tkToUnknown tk =
    TKUnknown
        { _tkuKey = keyPktTKKey (_tkPrimaryKey tk)
        , _tkuRevs = _tkRevs tk
        , _tkuUIDs = _tkUIDs tk
        , _tkuUAts = _tkUAts tk
        , _tkuSubs =
            map (\(kp, sigs) -> (keyPktToPkt kp, sigs)) (_tkSubs tk)
        }

someTKToUnknown :: SomeTK -> TKUnknown
someTKToUnknown (SomePublicTK tk) = tkToUnknown tk
someTKToUnknown (SomeSecretTK tk) = tkToUnknown tk

mkTKUnknown :: SomePKPayload -> Maybe SKAddendum -> TKUnknown
mkTKUnknown pkp maybeSka =
    TKUnknown
        { _tkuKey = (pkp, maybeSka)
        , _tkuRevs = []
        , _tkuUIDs = []
        , _tkuUAts = []
        , _tkuSubs = []
        }

fromPrimaryKeyPktToSomeTK :: Pkt -> Either String SomeTK
fromPrimaryKeyPktToSomeTK (PublicKeyPkt pkp) =
    Right
        ( SomePublicTK
            ( TK
                { _tkPrimaryKey = KeyPktPublicPrimary pkp
                , _tkRevs = []
                , _tkUIDs = []
                , _tkUAts = []
                , _tkSubs = []
                }
            )
        )
fromPrimaryKeyPktToSomeTK (SecretKeyPkt pkp ska) =
    Right
        ( SomeSecretTK
            ( TK
                { _tkPrimaryKey = KeyPktSecretPrimary pkp ska
                , _tkRevs = []
                , _tkUIDs = []
                , _tkUAts = []
                , _tkSubs = []
                }
            )
        )
fromPrimaryKeyPktToSomeTK pkt =
    Left
        ( "expected primary key packet, got packet tag "
            ++ show (pktTag pkt)
        )

someTKToPublicTK :: SomeTK -> Maybe (TK 'PublicTK)
someTKToPublicTK (SomePublicTK tk) = Just tk
someTKToPublicTK (SomeSecretTK _) = Nothing

someTKToSecretTK :: SomeTK -> Maybe (TK 'SecretTK)
someTKToSecretTK (SomeSecretTK tk) = Just tk
someTKToSecretTK (SomePublicTK _) = Nothing

someTKToPublicViewTK :: SomeTK -> TK 'PublicTK
someTKToPublicViewTK (SomePublicTK tk) = tk
someTKToPublicViewTK (SomeSecretTK tk) = publicViewTK tk

publicViewTK :: TK 'SecretTK -> TK 'PublicTK
publicViewTK tk =
    TK
        { _tkPrimaryKey = keyPktToPublicView (_tkPrimaryKey tk)
        , _tkRevs = _tkRevs tk
        , _tkUIDs = _tkUIDs tk
        , _tkUAts = _tkUAts tk
        , _tkSubs =
            map (\(kp, sigs) -> (keyPktToPublicView kp, sigs)) (_tkSubs tk)
        }

tkSomeSubs :: TK k -> [SomeKeyPkt]
tkSomeSubs tk =
    let xs = view (to _tkSubs) tk
     in xs ^.. (folded . _1 . to SomeKeyPkt)

fromUnknownToTKEither
    :: TKUnknown -> Either TKConversionError SomeTK
fromUnknownToTKEither tk =
    case _tkuKey tk of
        (pkp, Nothing) -> do
            subs <- traverse liftPublicSubkey (_tkuSubs tk)
            let typed :: TK 'PublicTK
                typed =
                    TK
                        { _tkPrimaryKey = KeyPktPublicPrimary pkp
                        , _tkRevs = _tkuRevs tk
                        , _tkUIDs = _tkuUIDs tk
                        , _tkUAts = _tkuUAts tk
                        , _tkSubs = subs
                        }
            Right
                (SomePublicTK typed)
        (pkp, Just ska) -> do
            subs <- traverse liftSecretSubkey (_tkuSubs tk)
            let typed :: TK 'SecretTK
                typed =
                    TK
                        { _tkPrimaryKey = KeyPktSecretPrimary pkp ska
                        , _tkRevs = _tkuRevs tk
                        , _tkUIDs = _tkuUIDs tk
                        , _tkUAts = _tkuUAts tk
                        , _tkSubs = subs
                        }
            Right
                (SomeSecretTK typed)
  where
    liftPublicSubkey
        :: (Pkt, [SignaturePayload])
        -> Either TKConversionError (KeyPkt 'PublicPkt, [SignaturePayload])
    liftPublicSubkey (pkt, sigs) =
        case pktToPublicKeyPkt pkt of
            Just keyPkt
                | keyPktRole keyPkt == KeyPktSubkey ->
                    Right (keyPkt, sigs)
                | otherwise ->
                    Left PublicSubkeyHasPrimaryRole
            Nothing ->
                Left (ExpectedPublicSubkeyPacket (pktTag pkt))

    liftSecretSubkey
        :: (Pkt, [SignaturePayload])
        -> Either TKConversionError (KeyPkt 'SecretPkt, [SignaturePayload])
    liftSecretSubkey (pkt, sigs) =
        case pktToSecretKeyPkt pkt of
            Just keyPkt
                | keyPktRole keyPkt == KeyPktSubkey ->
                    Right (keyPkt, sigs)
                | otherwise ->
                    Left SecretSubkeyHasPrimaryRole
            Nothing ->
                Left (ExpectedSecretSubkeyPacket (pktTag pkt))

fromUnknownToTK :: TKUnknown -> Either String SomeTK
fromUnknownToTK = first renderTKConversionError . fromUnknownToTKEither

instance Semigroup TKUnknown where
    a <> b =
        TKUnknown
            (_tkuKey a)
            (nub . sort $ _tkuRevs a ++ _tkuRevs b)
            ((kvmerge `on` _tkuUIDs) a b)
            ((kvmerge `on` _tkuUAts) a b)
            ((ukvmerge `on` _tkuSubs) a b)
      where
        kvmerge x y =
            Map.toList (Map.unionWith nsa (Map.fromList x) (Map.fromList y))
        ukvmerge x y =
            HashMap.toList
                (HashMap.unionWith nsa (HashMap.fromList x) (HashMap.fromList y))
        nsa x y = nub . sort $ x ++ y

instance Semigroup (TK k) where
    a <> b =
        TK
            (_tkPrimaryKey a)
            (nub . sort $ _tkRevs a ++ _tkRevs b)
            ((kvmerge `on` _tkUIDs) a b)
            ((kvmerge `on` _tkUAts) a b)
            ((ukvmerge `on` _tkSubs) a b)
      where
        kvmerge x y =
            Map.toList (Map.unionWith nsa (Map.fromList x) (Map.fromList y))
        ukvmerge x y =
            HashMap.toList
                (HashMap.unionWith nsa (HashMap.fromList x) (HashMap.fromList y))
        nsa x y = nub . sort $ x ++ y

instance Semigroup SomeTK where
    SomePublicTK a <> SomePublicTK b = SomePublicTK (a <> b)
    SomeSecretTK a <> SomeSecretTK b = SomeSecretTK (a <> b)
    a <> b =
        error
            ( "Semigroup SomeTK: cannot merge public and secret transferable keys"
                ++ " ("
                ++ show a
                ++ " <> "
                ++ show b
                ++ ")"
            )

data TKWithWireRep
    = TKWithWireRep
    { _tkWireRepRefs :: WireRepRefs
    , _tkWireRepRange :: Maybe ByteRange
    , _tkPackets :: [PktWithWireRep]
    , _tkValue :: TKUnknown
    }
    deriving (Data, Eq, Generic, Ord, Show, Typeable)

data PacketRefId
    = PacketRefId
    { _packetRefWireRepRef :: WireRepRef
    , _packetRefIndex :: Int
    }
    deriving (Data, Eq, Generic, Ord, Show, Typeable)

data SignatureWithWireRef
    = SignatureWithWireRef
    { _signatureWithWireRefValue :: SignaturePayload
    , _signatureWithWireRefRef :: PacketRefId
    }
    deriving (Data, Eq, Generic, Ord, Show, Typeable)

data UIDWithWireRefs
    = UIDWithWireRefs
    { _uidWithWireRefsValue :: Text
    , _uidWithWireRefsRef :: PacketRefId
    , _uidWithWireRefsSignatures :: [SignatureWithWireRef]
    }
    deriving (Data, Eq, Generic, Ord, Show, Typeable)

data UATWithWireRefs
    = UATWithWireRefs
    { _uatWithWireRefsValue :: [UserAttrSubPacket]
    , _uatWithWireRefsRef :: PacketRefId
    , _uatWithWireRefsSignatures :: [SignatureWithWireRef]
    }
    deriving (Data, Eq, Generic, Ord, Show, Typeable)

data SubkeyWithWireRefs
    = SubkeyWithWireRefs
    { _subkeyWithWireRefsValue :: Pkt
    , _subkeyWithWireRefsRef :: PacketRefId
    , _subkeyWithWireRefsSignatures :: [SignatureWithWireRef]
    }
    deriving (Data, Eq, Generic, Ord, Show, Typeable)

data TKStructuredWithWireRep
    = TKStructuredWithWireRep
    { _tkStructuredWireRepRefs :: WireRepRefs
    , _tkStructuredWireRepRange :: Maybe ByteRange
    , _tkStructuredPrimaryKey :: (SomePKPayload, Maybe SKAddendum)
    , _tkStructuredPrimaryKeyRef :: PacketRefId
    , _tkStructuredDirectSignatures :: [SignatureWithWireRef]
    , _tkStructuredUIDs :: [UIDWithWireRefs]
    , _tkStructuredUAts :: [UATWithWireRefs]
    , _tkStructuredSubkeys :: [SubkeyWithWireRefs]
    , _tkStructuredPacketRefs :: [PktWithWireRep]
    }
    deriving (Data, Eq, Generic, Ord, Show, Typeable)

data CanonicalizeTKWithWireRepError
    = CanonicalizeStructuringError String
    | CanonicalizeMissingPacketRef PacketRefId
    deriving (Data, Eq, Generic, Ord, Show, Typeable)

instance Ord TKUnknown where
    -- TKUnknown ordering is identity-oriented: the primary key packet defines key identity,
    -- while revocations, UIDs, and subkeys are mergeable metadata.
    compare = comparing _tkuKey

wireRepOfTK :: TKWithWireRep -> WireRepRef
wireRepOfTK = NE.head . _tkWireRepRefs

wireRepsOfTK :: TKWithWireRep -> WireRepRefs
wireRepsOfTK = _tkWireRepRefs

packetRefsOfTK :: TKWithWireRep -> [PktWithWireRep]
packetRefsOfTK = _tkPackets

packetRefIdOf :: PktWithWireRep -> PacketRefId
packetRefIdOf pkt = PacketRefId (_pktWireRepRef pkt) (_pktIndex pkt)

lookupPacketRef
    :: TKStructuredWithWireRep -> PacketRefId -> Maybe PktWithWireRep
lookupPacketRef structured target =
    find
        ((== target) . packetRefIdOf)
        (_tkStructuredPacketRefs structured)

packetWireBytesForRef
    :: TKStructuredWithWireRep
    -> PacketRefId
    -> Either CanonicalizeTKWithWireRepError BL.ByteString
packetWireBytesForRef structured refId =
    case lookupPacketRef structured refId of
        Just pkt -> Right (pkt ^. pktWireRep . pktRaw)
        Nothing -> Left (CanonicalizeMissingPacketRef refId)

signatureWireSortKey
    :: TKStructuredWithWireRep
    -> SignatureWithWireRef
    -> Either
        CanonicalizeTKWithWireRepError
        (BL.ByteString, PacketRefId)
signatureWireSortKey structured sig =
    (\raw -> (raw, _signatureWithWireRefRef sig))
        <$> packetWireBytesForRef structured (_signatureWithWireRefRef sig)

uidWireSortKey
    :: TKStructuredWithWireRep
    -> UIDWithWireRefs
    -> Either
        CanonicalizeTKWithWireRepError
        (BL.ByteString, PacketRefId)
uidWireSortKey structured uid =
    (\raw -> (raw, _uidWithWireRefsRef uid))
        <$> packetWireBytesForRef structured (_uidWithWireRefsRef uid)

uatWireSortKey
    :: TKStructuredWithWireRep
    -> UATWithWireRefs
    -> Either
        CanonicalizeTKWithWireRepError
        (BL.ByteString, PacketRefId)
uatWireSortKey structured uat =
    (\raw -> (raw, _uatWithWireRefsRef uat))
        <$> packetWireBytesForRef structured (_uatWithWireRefsRef uat)

subkeyWireSortKey
    :: TKStructuredWithWireRep
    -> SubkeyWithWireRefs
    -> Either
        CanonicalizeTKWithWireRepError
        (BL.ByteString, PacketRefId)
subkeyWireSortKey structured sub =
    (\raw -> (raw, _subkeyWithWireRefsRef sub))
        <$> packetWireBytesForRef structured (_subkeyWithWireRefsRef sub)

compareSignatureWithWireRefCanonical
    :: TKStructuredWithWireRep
    -> SignatureWithWireRef
    -> SignatureWithWireRef
    -> Either CanonicalizeTKWithWireRepError Ordering
compareSignatureWithWireRefCanonical structured a b =
    compare
        <$> signatureWireSortKey structured a
        <*> signatureWireSortKey structured b

compareUIDWithWireRefsCanonical
    :: TKStructuredWithWireRep
    -> UIDWithWireRefs
    -> UIDWithWireRefs
    -> Either CanonicalizeTKWithWireRepError Ordering
compareUIDWithWireRefsCanonical structured a b =
    compare
        <$> uidWireSortKey structured a
        <*> uidWireSortKey structured b

compareUATWithWireRefsCanonical
    :: TKStructuredWithWireRep
    -> UATWithWireRefs
    -> UATWithWireRefs
    -> Either CanonicalizeTKWithWireRepError Ordering
compareUATWithWireRefsCanonical structured a b =
    compare
        <$> uatWireSortKey structured a
        <*> uatWireSortKey structured b

compareSubkeyWithWireRefsCanonical
    :: TKStructuredWithWireRep
    -> SubkeyWithWireRefs
    -> SubkeyWithWireRefs
    -> Either CanonicalizeTKWithWireRepError Ordering
compareSubkeyWithWireRefsCanonical structured a b =
    compare
        <$> subkeyWireSortKey structured a
        <*> subkeyWireSortKey structured b

sortCanonicalByKey
    :: Ord key
    => (a -> Either CanonicalizeTKWithWireRepError key)
    -> [a]
    -> Either CanonicalizeTKWithWireRepError [a]
sortCanonicalByKey keyFn xs =
    map snd . sortOn fst
        <$> traverse (\x -> (\k -> (k, x)) <$> keyFn x) xs

sortSignatureWithWireRefsCanonical
    :: TKStructuredWithWireRep
    -> [SignatureWithWireRef]
    -> Either CanonicalizeTKWithWireRepError [SignatureWithWireRef]
sortSignatureWithWireRefsCanonical structured =
    sortCanonicalByKey (signatureWireSortKey structured)

sortUIDWithWireRefsCanonical
    :: TKStructuredWithWireRep
    -> [UIDWithWireRefs]
    -> Either CanonicalizeTKWithWireRepError [UIDWithWireRefs]
sortUIDWithWireRefsCanonical structured uids = do
    normalized <-
        traverse
            ( \uid ->
                (\sigs -> uid {_uidWithWireRefsSignatures = sigs})
                    <$> sortSignatureWithWireRefsCanonical
                        structured
                        (_uidWithWireRefsSignatures uid)
            )
            uids
    sortCanonicalByKey (uidWireSortKey structured) normalized

sortUATWithWireRefsCanonical
    :: TKStructuredWithWireRep
    -> [UATWithWireRefs]
    -> Either CanonicalizeTKWithWireRepError [UATWithWireRefs]
sortUATWithWireRefsCanonical structured uats = do
    normalized <-
        traverse
            ( \uat ->
                (\sigs -> uat {_uatWithWireRefsSignatures = sigs})
                    <$> sortSignatureWithWireRefsCanonical
                        structured
                        (_uatWithWireRefsSignatures uat)
            )
            uats
    sortCanonicalByKey (uatWireSortKey structured) normalized

sortSubkeyWithWireRefsCanonical
    :: TKStructuredWithWireRep
    -> [SubkeyWithWireRefs]
    -> Either CanonicalizeTKWithWireRepError [SubkeyWithWireRefs]
sortSubkeyWithWireRefsCanonical structured subs = do
    normalized <-
        traverse
            ( \sub ->
                (\sigs -> sub {_subkeyWithWireRefsSignatures = sigs})
                    <$> sortSignatureWithWireRefsCanonical
                        structured
                        (_subkeyWithWireRefsSignatures sub)
            )
            subs
    sortCanonicalByKey (subkeyWireSortKey structured) normalized

canonicalizeTKStructuredWithWireRep
    :: TKStructuredWithWireRep
    -> Either CanonicalizeTKWithWireRepError TKUnknown
canonicalizeTKStructuredWithWireRep structured =
    buildTK
        <$> sortSignatureWithWireRefsCanonical
            structured
            (_tkStructuredDirectSignatures structured)
        <*> sortUIDWithWireRefsCanonical
            structured
            (_tkStructuredUIDs structured)
        <*> sortUATWithWireRefsCanonical
            structured
            (_tkStructuredUAts structured)
        <*> sortSubkeyWithWireRefsCanonical
            structured
            (_tkStructuredSubkeys structured)
  where
    buildTK directSigs uids uats subs =
        TKUnknown
            { _tkuKey = _tkStructuredPrimaryKey structured
            , _tkuRevs = map _signatureWithWireRefValue directSigs
            , _tkuUIDs =
                map
                    ( _uidWithWireRefsValue
                        &&& (map _signatureWithWireRefValue . _uidWithWireRefsSignatures)
                    )
                    uids
            , _tkuUAts =
                map
                    ( _uatWithWireRefsValue
                        &&& (map _signatureWithWireRefValue . _uatWithWireRefsSignatures)
                    )
                    uats
            , _tkuSubs =
                map
                    ( _subkeyWithWireRefsValue
                        &&& (map _signatureWithWireRefValue . _subkeyWithWireRefsSignatures)
                    )
                    subs
            }

canonicalizeTKWithWireRep
    :: TKWithWireRep -> Either CanonicalizeTKWithWireRepError TKUnknown
canonicalizeTKWithWireRep tk = do
    structured <-
        case toStructuredTKWithWireRep tk of
            Left err -> Left (CanonicalizeStructuringError err)
            Right s -> Right s
    canonicalizeTKStructuredWithWireRep structured

toStructuredTKWithWireRep
    :: TKWithWireRep -> Either String TKStructuredWithWireRep
toStructuredTKWithWireRep tkWithRefs = do
    let tk = _tkValue tkWithRefs
        refs = _tkPackets tkWithRefs
        (pkp, mska) = _tkuKey tk
        primaryPkt = someKeyPktToPkt (mkPrimaryKeyPkt pkp mska)
    zipper <- case zFromList refs of
        Just z -> Right z
        Nothing ->
            Left "no packet references available for TKUnknown structuring"
    (primaryRef, z1') <-
        consumePktZ "primary key packet" primaryPkt zipper
    -- Move past the primary key to process its signatures and following packets
    z1 <- case zMoveNext z1' of
        Just z -> Right z
        Nothing ->
            -- Primary key is the only packet; only valid if no revisions, UIDs, UATs, or subkeys
            if null (_tkuRevs tk)
                && null (_tkuUIDs tk)
                && null (_tkuUAts tk)
                && null (_tkuSubs tk)
                then Right z1'
                else
                    Left "missing signatures/UIDs/subkeys after primary key packet"
    (directSigs, z2) <-
        consumeSigsZ "direct-key signatures" (_tkuRevs tk) z1
    (uids, z3) <- consumeUIDsZ (_tkuUIDs tk) z2
    (uats, z4) <- consumeUATsZ (_tkuUAts tk) z3
    (subs, z5) <- consumeSubsZ (_tkuSubs tk) z4
    -- Check if there are trailing packets AFTER the current focus (not including it)
    case _zpAfter z5 of
        [] ->
            Right
                ( TKStructuredWithWireRep
                    (_tkWireRepRefs tkWithRefs)
                    (_tkWireRepRange tkWithRefs)
                    (_tkuKey tk)
                    (packetRefIdOf primaryRef)
                    directSigs
                    uids
                    uats
                    subs
                    (_tkPackets tkWithRefs)
                )
        (unexpected : _) ->
            Left
                ( "unexpected trailing packet reference at position "
                    ++ show (zPosition z5 + 1)
                    ++ " while structuring TKUnknown provenance (tag "
                    ++ show (pktTag (unexpected ^. pktWireRep . pktValue))
                    ++ ")"
                )
  where
    consumePktZ
        :: String
        -> Pkt
        -> PacketZipper PktWithWireRep
        -> Either String (PktWithWireRep, PacketZipper PktWithWireRep)
    consumePktZ context expected z
        | (extract z) ^. pktWireRep . pktValue == expected =
            Right (extract z, z)
        | otherwise =
            Left
                ( "packet/reference mismatch for "
                    ++ context
                    ++ " at position "
                    ++ show (zPosition z)
                    ++ ": expected tag "
                    ++ show (pktTag expected)
                    ++ ", got tag "
                    ++ show (pktTag ((extract z) ^. pktWireRep . pktValue))
                )

    tryMoveNext
        :: String
        -> PacketZipper PktWithWireRep
        -> Bool
        -> Either String (PacketZipper PktWithWireRep)
    tryMoveNext errorPrefix z isLastItem =
        case zMoveNext z of
            Just z' -> Right z'
            Nothing ->
                if isLastItem
                    then Right z
                    else
                        Left
                            ( errorPrefix
                                ++ " at position "
                                ++ show (zPosition z)
                            )

    consumeSigsZ
        :: String
        -> [SignaturePayload]
        -> PacketZipper PktWithWireRep
        -> Either
            String
            ([SignatureWithWireRef], PacketZipper PktWithWireRep)
    consumeSigsZ context sigs z = go [] sigs z
      where
        go acc [] zipper = Right (reverse acc, zipper)
        go acc (sig : sigRest) zipper = do
            (sigPkt, z') <- consumePktZ context (SignaturePkt sig) zipper
            z'' <- tryMoveNext "missing signature packet" z' (null sigRest)
            go
                (SignatureWithWireRef sig (packetRefIdOf sigPkt) : acc)
                sigRest
                z''

    consumeUIDsZ
        :: [(Text, [SignaturePayload])]
        -> PacketZipper PktWithWireRep
        -> Either String ([UIDWithWireRefs], PacketZipper PktWithWireRep)
    consumeUIDsZ [] z = Right ([], z)
    consumeUIDsZ ((uid, sigs) : rest) z = do
        (uidPkt, z1) <- consumePktZ "UID packet" (UserIdPkt uid) z
        z2 <- tryMoveNext "missing UID" z1 (null rest && null sigs)
        (uidSigs, z3) <- consumeSigsZ "UID signature" sigs z2
        (tailUIDs, z4) <- consumeUIDsZ rest z3
        Right
            ( UIDWithWireRefs uid (packetRefIdOf uidPkt) uidSigs : tailUIDs
            , z4
            )

    consumeUATsZ
        :: [([UserAttrSubPacket], [SignaturePayload])]
        -> PacketZipper PktWithWireRep
        -> Either String ([UATWithWireRefs], PacketZipper PktWithWireRep)
    consumeUATsZ [] z = Right ([], z)
    consumeUATsZ ((uat, sigs) : rest) z = do
        (uatPkt, z1) <- consumePktZ "UAT packet" (UserAttributePkt uat) z
        z2 <- tryMoveNext "missing UAT" z1 (null rest && null sigs)
        (uatSigs, z3) <- consumeSigsZ "UAT signature" sigs z2
        (tailUats, z4) <- consumeUATsZ rest z3
        Right
            ( UATWithWireRefs uat (packetRefIdOf uatPkt) uatSigs : tailUats
            , z4
            )

    consumeSubsZ
        :: [(Pkt, [SignaturePayload])]
        -> PacketZipper PktWithWireRep
        -> Either String ([SubkeyWithWireRefs], PacketZipper PktWithWireRep)
    consumeSubsZ [] z = Right ([], z)
    consumeSubsZ ((subPkt, sigs) : rest) z = do
        (subRef, z1) <- consumePktZ "subkey packet" subPkt z
        z2 <- tryMoveNext "missing subkey" z1 (null rest && null sigs)
        (subSigs, z3) <- consumeSigsZ "subkey signature" sigs z2
        (tailSubs, z4) <- consumeSubsZ rest z3
        Right
            ( SubkeyWithWireRefs subPkt (packetRefIdOf subRef) subSigs
                : tailSubs
            , z4
            )

tksFromWireRep
    :: WireRepRef -> [TKWithWireRep] -> [TKWithWireRep]
tksFromWireRep src = filter (elem src . NE.toList . wireRepsOfTK)

tksContainingPacket
    :: PktWithWireRep -> [TKWithWireRep] -> [TKWithWireRep]
tksContainingPacket pkt = filter (elem pkt . packetRefsOfTK)

$(ATH.deriveToJSON ATH.defaultOptions ''TKUnknown)

type KeyringIxs = '[EightOctetKeyId, Fingerprint, Text]

-- | Kinded keyrings: homogeneous collections of public or secret TKs
type PublicKeyring = IxSet KeyringIxs (TK 'PublicTK)

type SecretKeyring = IxSet KeyringIxs (TK 'SecretTK)

-- | Parameterized kinded keyring for generic operations
type family KeyringOf (k :: TKKind) :: Type where
    KeyringOf 'PublicTK = PublicKeyring
    KeyringOf 'SecretTK = SecretKeyring

$(makeLenses ''TKUnknown)
$(makeLenses ''TK)
$(makeLenses ''TKWithWireRep)
$(makeLenses ''PacketRefId)
$(makeLenses ''PacketZipper)
$(makeLenses ''SignatureWithWireRef)
$(makeLenses ''UIDWithWireRefs)
$(makeLenses ''UATWithWireRefs)
$(makeLenses ''SubkeyWithWireRefs)
$(makeLenses ''TKStructuredWithWireRep)
