{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE PartialTypeSignatures      #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE TupleSections              #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE TypeFamilyDependencies     #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE UndecidableInstances       #-}
module OnionTransport
    ( parseOnionAddr
    , encodeOnionAddr
    , forwardOnions
    , OnionToOwner(..)
    , OnionMessage(..)
    , DataToRoute(..)
    , AnnounceResponse(..)
    , AnnounceRequest(..)
    , Forwarding(..)
    , ReturnPath(..)
    , OnionRequest(..)
    , OnionResponse(..)
    , Addressed(..)
    , UDPTransport
    , KeyRecord(..)
    , encrypt
    , decrypt
    ) where

import Network.Address (fromSockAddr,toSockAddr,setPort,either4or6,sockAddrPort)
import Network.QueryResponse
import ToxCrypto hiding (encrypt,decrypt)
import ToxAddress
import qualified ToxCrypto
import DHTTransport (NodeInfo(..),NodeId(..),SendNodes,nodeInfo,DHTPublicKey)

import Control.Arrow
import Control.Concurrent.STM
import qualified Data.ByteString as B
         ;import Data.ByteString (ByteString)
import Data.Coerce
import Data.Functor.Contravariant
import Data.Functor.Identity
import Data.IP
import Data.Maybe
import Data.Monoid
import Data.Serialize            as S
import Data.Typeable
import Data.Word
import GHC.TypeLits
import Network.Socket
import GHC.Generics ()
import Data.Type.Equality

type HandleLo a = Maybe (Either String (ByteString, SockAddr)) -> IO a

type UDPTransport = Transport String SockAddr ByteString


getOnionAssym :: Get (Assym (Encrypted DataToRoute))
getOnionAssym = getAliasedAssym

putOnionAssym :: Serialize a => Word8 -> Put -> Assym a -> Put
putOnionAssym typ p a = put typ >> p >> putAliasedAssym a

data OnionMessage (f :: * -> *)
    = OnionAnnounce (Assym (f (AnnounceRequest,Nonce8)))
    | OnionAnnounceResponse Nonce8 Nonce24 (f AnnounceResponse)
    | OnionToRoute PublicKey (Assym (f DataToRoute)) -- destination key, aliased Assym
    | OnionToRouteResponse (Assym (f DataToRoute))

deriving instance ( Show (f (AnnounceRequest, Nonce8))
                  , Show (f AnnounceResponse)
                  , Show (f DataToRoute)
                  ) => Show (OnionMessage f)

data OnionToOwner = OnionToOwner NodeInfo (ReturnPath N3)
                  | OnionToMe SockAddr -- SockAddr is immediate peer in route
    deriving Show

onionKey :: OnionToOwner -> Maybe PublicKey
onionKey (OnionToOwner ni _) = Just $ id2key (nodeId ni)
onionKey _                   = Nothing

instance Sized (OnionMessage Encrypted) where
    size = VarSize $ \case
        OnionAnnounce a -> case size of ConstSize n -> n + 1
                                        VarSize f   -> f a + 1
        OnionAnnounceResponse n8 n24 x -> case size of ConstSize n -> n + 33
                                                       VarSize f   -> f x + 33
        OnionToRoute pubkey a -> case size of ConstSize n -> n + 33
                                              VarSize f   -> f a + 33
        OnionToRouteResponse a -> case size of ConstSize n -> n + 1
                                               VarSize f   -> f a + 1

instance Serialize (OnionMessage Encrypted) where
    get = do
        typ <- get
        case typ :: Word8 of
            0x83 -> OnionAnnounce <$> getAliasedAssym
            0x85 -> OnionToRoute <$> getPublicKey <*> getAliasedAssym
            0x84 -> getOnionReply typ
            0x86 -> getOnionReply typ
            t    -> fail $ "Unknown onion payload: " ++ show t
    put (OnionAnnounce a)                = putWord8 0x83 >> putAliasedAssym a
    put (OnionToRoute k a)               = putWord8 0x85 >> putPublicKey k >> putAliasedAssym a
    put (OnionAnnounceResponse n8 n24 x) = putWord8 0x84 >> put n8 >> put n24 >> put x
    put (OnionToRouteResponse a)         = putWord8 0x86 >> putAliasedAssym a

onionToOwner :: Assym a -> ReturnPath N3 -> SockAddr -> Either String OnionToOwner
onionToOwner assym ret3 saddr = do
    ni <- nodeInfo (key2id $ senderKey assym) saddr
    return $ OnionToOwner ni ret3
-- data CookieAddress = WithoutCookie NodeInfo | CookieAddress Cookie SockAddr


onion :: Sized msg =>
               ByteString
               -> SockAddr
               -> Get (Assym (Encrypted msg) -> t)
               -> Either String (t, OnionToOwner)
onion bs saddr getf = do (f,(assym,ret3)) <- runGet ((,) <$> getf <*> getOnionRequest) bs
                         oaddr <- onionToOwner assym ret3 saddr
                         return (f assym, oaddr)


parseOnionAddr :: (ByteString, SockAddr) -> Either (OnionMessage Encrypted,OnionToOwner) (ByteString,SockAddr)
parseOnionAddr (msg,saddr)
    | Just (typ,bs) <- B.uncons msg
    , let right = Right (msg,saddr)
          query = either (const right) Left
          response = either (const right) (Left . (, OnionToMe saddr))
    = case typ of
    0x83 -> query $ onion bs saddr (pure OnionAnnounce)            -- Announce Request
    0x85 -> query $ onion bs saddr (OnionToRoute <$> getPublicKey) -- Onion Data Request
    0x84 -> response $ runGet (getOnionReply 0x84) bs              -- Announce Response
    0x86 -> response $ runGet (getOnionReply 0x86) bs              -- Onion Data Response
    _ -> right

getOnionReply :: Word8 -> Get (OnionMessage Encrypted)
getOnionReply 0x84 = OnionAnnounceResponse <$> get <*> get <*> get
getOnionReply 0x86 = OnionToRouteResponse <$> getOnionAssym

putOnionMsg :: OnionMessage Encrypted -> Put
putOnionMsg (OnionAnnounce                a) = putOnionAssym 0x83 (return ()) a
putOnionMsg (OnionToRoute pubkey          a) = putOnionAssym 0x85 (putPublicKey pubkey) a
putOnionMsg (OnionAnnounceResponse n8 n24 x) = put (0x84 :: Word8) >> put n8 >> put n24 >> put x
putOnionMsg (OnionToRouteResponse         a) = putOnionAssym 0x86 (return ()) a

encodeOnionAddr :: (OnionMessage Encrypted,OnionToOwner) -> (ByteString, SockAddr)
encodeOnionAddr (msg,addr) = ( runPut (putOnionMsg msg >> putpath), saddr )
 where
    (saddr,putpath) | OnionToOwner ni p <- addr = (nodeAddr ni, put p)
                    | OnionToMe a       <- addr = (a, return ())

forwardOnions :: TransportCrypto -> UDPTransport -> UDPTransport -- HandleHi a -> IO a -> HandleLo a
forwardOnions crypto udp = udp { awaitMessage = await' }
 where
    -- forMe :: HandleHi
    -- forThem :: handleLo
    await' :: HandleLo a -> IO a
    await' forThem = awaitMessage udp $ \case
        m@(Just (Right (bs,saddr))) -> case B.head bs of
            0x80 -> forward forThem bs $ handleOnionRequest (Proxy :: Proxy  N0) crypto saddr udp (await' forThem)
            0x81 -> forward forThem bs $ handleOnionRequest (Proxy :: Proxy  N1) crypto saddr udp (await' forThem)
            0x82 -> forward forThem bs $ handleOnionRequest (Proxy :: Proxy  N2) crypto saddr udp (await' forThem)
            0x8c -> forward forThem bs $ handleOnionResponse (Proxy :: Proxy N3) crypto saddr udp (await' forThem)
            0x8d -> forward forThem bs $ handleOnionResponse (Proxy :: Proxy N2) crypto saddr udp (await' forThem)
            0x8e -> forward forThem bs $ handleOnionResponse (Proxy :: Proxy N1) crypto saddr udp (await' forThem)
            _    -> forThem m
        m -> forThem m

forward :: forall c b b1. Serialize b =>
              (Maybe (Either String b1) -> c) -> ByteString -> (b -> c) -> c
forward forMe bs f = either (forMe . Just . Left) f $ decode $ B.tail bs

class SumToThree a b

instance SumToThree N0 N3
instance SumToThree (S a) b => SumToThree a (S b)

class ( Serialize (ReturnPath n)
      , Serialize (ReturnPath (S n))
      , Serialize (Forwarding (ThreeMinus (S n)) (OnionMessage Encrypted))
      , ThreeMinus n ~ S (ThreeMinus (S n))
      ) => LessThanThree n

instance LessThanThree N0
instance LessThanThree N1
instance LessThanThree N2

type family ThreeMinus n = r | r -> n where
    ThreeMinus N3 = N0
    ThreeMinus N2 = N1
    ThreeMinus N1 = N2
    ThreeMinus N0 = N3

-- n = 0, 1, 2
data OnionRequest n = OnionRequest
    { onionNonce    :: Nonce24
    , onionForward  :: Forwarding (ThreeMinus n) (OnionMessage Encrypted)
    , pathFromOwner :: ReturnPath n
    }

instance ( Serialize (Forwarding (ThreeMinus n) (OnionMessage Encrypted))
         , Serialize (ReturnPath n)
         ) => Serialize (OnionRequest n) where
    get = OnionRequest <$> get <*> get <*> get
    put (OnionRequest n f p) = put n >> put f >> put p

-- getRequest :: _
-- getRequest = OnionRequest <$> get <*> get <*> get

-- n = 1, 2, 3
-- Attributed (Encrypted (

data OnionResponse n = OnionResponse
    { pathToOwner :: ReturnPath n
    , msgToOwner  :: OnionMessage Encrypted
    }

instance ( Serialize (ReturnPath n) ) => Serialize (OnionResponse n) where
    get = OnionResponse <$> get <*> (get >>= getOnionReply)
    put (OnionResponse p m) = put p >> putOnionMsg m


data Addressed a = Addressed { sockAddr :: SockAddr, unaddressed :: a }
 deriving (Eq,Show)

instance Sized a => Sized (Addressed a) where
    size = case size :: Size a of
            ConstSize n -> ConstSize     $ 1{-family-} + 16{-ip-} + 2{-port-} + n
            VarSize f   -> VarSize $ \x -> 1{-family-} + 16{-ip-} + 2{-port-} + f (unaddressed x)

getForwardAddr :: S.Get SockAddr
getForwardAddr = do
    addrfam <- S.get :: S.Get Word8
    ip <- getIP addrfam
    case ip of IPv4 _ -> S.skip 12 -- compliant peers would zero-fill this.
               IPv6 _ -> return ()
    port <- S.get :: S.Get PortNumber
    return $ setPort port $ toSockAddr ip


putForwardAddr :: SockAddr -> S.Put
putForwardAddr saddr = fromMaybe (return $ error "unsupported SockAddr family") $ do
    port <- sockAddrPort saddr
    ip <- fromSockAddr $ either id id $ either4or6 saddr
    return $ do
        case ip of
            IPv4 ip4 -> S.put (0x02 :: Word8) >> S.put ip4 >> S.putByteString (B.replicate 12 0)
            IPv6 ip6 -> S.put (0x0a :: Word8) >> S.put ip6
        S.put port

instance Serialize a => Serialize (Addressed a) where
    get = Addressed <$> getForwardAddr <*> get
    put (Addressed addr x) = putForwardAddr addr >> put x

data N0
data S n
type N1 = S N0
type N2 = S N1
type N3 = S N2

type family PeanoNat p where
    PeanoNat N0    = 0
    PeanoNat (S n) = 1 + PeanoNat n

data ReturnPath n where
    NoReturnPath :: ReturnPath N0
    ReturnPath :: Nonce24 -> Encrypted (Addressed (ReturnPath n)) -> ReturnPath (S n)

-- Size: 59 = 1(family) + 16(ip) + 2(port) +16(mac) + 24(nonce)
instance Sized (ReturnPath N0) where size = ConstSize 0
instance Sized (ReturnPath n) => Sized (ReturnPath (S n)) where
    size = ConstSize 59 <> contramap (\x -> let _ = x :: ReturnPath (S n)
                                            in error "non-constant ReturnPath size")
                                     (size :: Size (ReturnPath n))

{-
instance KnownNat (PeanoNat n) => Sized (ReturnPath n) where
    size = ConstSize $ 59 * fromIntegral (natVal (Proxy :: Proxy (PeanoNat n)))
-}

instance Serialize (ReturnPath N0) where get              = pure NoReturnPath
                                         put NoReturnPath = pure ()

instance Serialize (ReturnPath N1) where
    get = ReturnPath <$> get <*> get
    put (ReturnPath n24 p) = put n24 >> put p

instance (Sized (ReturnPath n), Serialize (ReturnPath n)) => Serialize (ReturnPath (S (S n))) where
    get = ReturnPath <$> get <*> get
    put (ReturnPath n24 p) = put n24 >> put p


{-
-- This doesn't work because it tried to infer it for (0 - 1)
instance (Serialize (Encrypted (Addressed (ReturnPath (n - 1))))) => Serialize (ReturnPath n) where
    get = ReturnPath <$> get <*> get
    put (ReturnPath n24 p) = put n24 >> put p
-}

instance KnownNat (PeanoNat n) => Show (ReturnPath n) where
    show rpath = "ReturnPath" ++ show (natVal (Proxy :: Proxy (PeanoNat n)))


-- instance KnownNat n => Serialize (ReturnPath n) where
--     -- Size: 59 = 1(family) + 16(ip) + 2(port) +16(mac) + 24(nonce)
--     get = ReturnPath <$> getBytes ( 59 * (fromIntegral $ natVal $ Proxy @n) )
--     put (ReturnPath bs) = putByteString bs


data Forwarding n msg where
    NotForwarded :: msg -> Forwarding N0 msg
    Forwarding :: Assym (Encrypted (Addressed (Forwarding n msg))) -> Forwarding (S n) msg

instance Sized msg => Sized (Forwarding N0 msg)
   where size = case size :: Size msg of
                 ConstSize n -> ConstSize n
                 VarSize f -> VarSize $ \(NotForwarded x) -> f x

instance Sized (Forwarding n msg) => Sized (Forwarding (S n) msg)
   where size = case size :: Size (Assym (Encrypted (Addressed (Forwarding n msg)))) of
                 ConstSize n -> ConstSize n
                 VarSize f -> VarSize $ \(Forwarding a) -> f a

instance Serialize msg => Serialize (Forwarding N0 msg) where
    get = NotForwarded <$> get
    put (NotForwarded msg) = put msg

instance (Serialize (Encrypted (Addressed (Forwarding n msg)))) => Serialize (Forwarding (S n) msg) where
    get = Forwarding <$> getAliasedAssym
    put (Forwarding x) = putAliasedAssym x

handleOnionRequest :: LessThanThree n => proxy n -> TransportCrypto -> SockAddr -> UDPTransport -> IO a -> OnionRequest n -> IO a
handleOnionRequest proxy crypto saddr udp kont (OnionRequest nonce msg rpath) = do
    (sym, snonce) <- atomically ( (,) <$> transportSymmetric crypto
                                      <*> transportNewNonce crypto )
    case peelOnion crypto msg of
        Left e -> do
            -- todo report encryption error
            kont
        Right (Addressed dst msg') -> do
            sendMessage udp dst (S.encode $ OnionRequest nonce msg' $ wrapSymmetric sym snonce saddr rpath)
            kont

wrapSymmetric :: Serialize (ReturnPath n) =>
                 SymmetricKey -> Nonce24 -> SockAddr -> ReturnPath n -> ReturnPath (S n)
wrapSymmetric sym n saddr rpath = ReturnPath n $ encryptSymmetric sym n (encodePlain $ Addressed saddr rpath)

peelSymmetric :: Serialize (Addressed (ReturnPath n))
                 => SymmetricKey -> ReturnPath (S n) -> Either String (Addressed (ReturnPath n))
peelSymmetric sym (ReturnPath nonce e) = decryptSymmetric sym nonce e >>= decodePlain


peelOnion :: Serialize (Addressed (Forwarding n t))
              => TransportCrypto
              -> Forwarding (S n) t
              -> Either String (Addressed (Forwarding n t))
peelOnion crypto (Forwarding fwd) =
    fmap runIdentity $ uncomposed $ decryptMessage crypto (assymNonce fwd) (Right fwd)

handleOnionResponse :: Serialize (ReturnPath n) => proxy (S n) -> TransportCrypto -> SockAddr -> UDPTransport -> IO a -> OnionResponse (S n) -> IO a
handleOnionResponse proxy crypto saddr udp kont (OnionResponse path msg) = do
    sym <- atomically $ transportSymmetric crypto
    case peelSymmetric sym path of
        Left e -> do
            -- todo report encryption error
            kont
        Right (Addressed dst path') -> do
            sendMessage udp dst (S.encode $ OnionResponse path' msg)
            kont


data AnnounceRequest = AnnounceRequest
    { announcePingId  :: Nonce32 -- Ping ID
    , announceSeeking :: NodeId  -- Public key we are searching for
    , announceKey     :: NodeId  -- Public key that we want those sending back data packets to use
    }

instance Sized AnnounceRequest where size = ConstSize (32*3)

instance S.Serialize AnnounceRequest where
    get = AnnounceRequest <$> S.get <*> S.get <*> S.get
    put (AnnounceRequest p s k) = S.put (p,s,k)

getOnionRequest :: Sized msg => Get (Assym (Encrypted msg), ReturnPath N3)
getOnionRequest = do
    -- Assumes return path is constant size so that we can isolate
    -- the variable-sized prefix.
    cnt <- remaining
    a <- isolate (case size :: Size (ReturnPath N3) of ConstSize n -> cnt - n)
                 getAliasedAssym
    path <- get
    return (a,path)

data KeyRecord = NotStored    Nonce32
               | SendBackKey  PublicKey
               | Acknowledged Nonce32

instance Sized KeyRecord where size = ConstSize 33

instance S.Serialize KeyRecord where
    get = do
        is_stored <- S.get :: S.Get Word8
        case is_stored of
            1 -> SendBackKey  <$> getPublicKey
            2 -> Acknowledged <$> S.get
            _ -> NotStored    <$> S.get
    put (NotStored n32)    = S.put (0 :: Word8) >> S.put n32
    put (SendBackKey key)  = S.put (1 :: Word8) >> putPublicKey key
    put (Acknowledged n32) = S.put (2 :: Word8) >> S.put n32

data AnnounceResponse = AnnounceResponse
    { is_stored     :: KeyRecord
    , announceNodes :: SendNodes
    }

instance Sized AnnounceResponse where
    size = contramap is_stored size <> contramap announceNodes size

instance S.Serialize AnnounceResponse where
    get = AnnounceResponse <$> S.get <*> S.get
    put (AnnounceResponse st ns) = S.put st >> S.put ns

data DataToRoute = DataToRoute
    { dataFromKey :: PublicKey -- Real public key of sender
    , dataToRoute :: Encrypted OnionData -- (Word8,ByteString) -- DHTPK 0x9c
    }

instance Sized DataToRoute where
    size = ConstSize 32 <> contramap dataToRoute size

instance Serialize DataToRoute where
    get = DataToRoute <$> getPublicKey <*> get
    put (DataToRoute k dta) = putPublicKey k >> put dta

data OnionData = OnionDHTPublicKey DHTPublicKey -- 0x9c

instance Sized OnionData where
    size = VarSize $ \(OnionDHTPublicKey dhtpk) -> case size of
            ConstSize n -> n -- Override because OnionData probably
                             -- should be treated as variable sized.
            VarSize f   -> f dhtpk

encrypt :: TransportCrypto -> OnionMessage Identity -> OnionToOwner -> (OnionMessage Encrypted, OnionToOwner)
encrypt crypto msg rpath = ( transcode (encryptMessage crypto okey) msg
                           , rpath)
    where
        -- The OnionToMe case shouldn't happen, but we'll use our own public
        -- key in this situation.
        okey = fromMaybe (transportPublic crypto) $ onionKey rpath

encryptMessage :: Serialize a =>
                  TransportCrypto -> PublicKey -> Nonce24 -> Either (Identity a) (Assym (Identity a)) -> Encrypted a
encryptMessage crypto destKey n (Right a) = ToxCrypto.encrypt secret plain
 where
    secret = computeSharedSecret (transportSecret crypto) destKey n
    plain  = encodePlain $ runIdentity $ assymData a
encryptMessage crypto destKey n (Left x) = ToxCrypto.encrypt secret plain
 where
    secret = computeSharedSecret (transportSecret crypto) destKey n
    plain  = encodePlain $ runIdentity $ x

decrypt :: TransportCrypto -> OnionMessage Encrypted -> OnionToOwner -> Either String (OnionMessage Identity, OnionToOwner)
decrypt crypto msg addr = (, addr) <$> (sequenceMessage $ transcode (decryptMessage crypto) msg)

decryptMessage :: Serialize x =>
                  TransportCrypto
                  -> Nonce24
                  -> Either (Encrypted x) (Assym (Encrypted x))
                  -> (Either String ∘ Identity) x
decryptMessage crypto n (Right assymE) = plain $ ToxCrypto.decrypt secret e
 where
    secret = computeSharedSecret (transportSecret crypto) (senderKey assymE) n
    e      = assymData assymE
    plain  = Composed . fmap Identity . (>>= decodePlain)
decryptMessage crypto n (Left e) = _todo -- OnionAnnounceResponse has no sender key


sequenceMessage :: Applicative m => OnionMessage (m ∘ f) -> m (OnionMessage f)
sequenceMessage (OnionAnnounce a)        = fmap OnionAnnounce        $ sequenceA $ fmap uncomposed a
sequenceMessage (OnionToRoute pub a)     = fmap (OnionToRoute pub)   $ sequenceA $ fmap uncomposed a
sequenceMessage (OnionToRouteResponse a) = fmap OnionToRouteResponse $ sequenceA $ fmap uncomposed a
sequenceMessage (OnionAnnounceResponse n8 n24 dta) = OnionAnnounceResponse n8 n24 <$> uncomposed dta

transcode :: forall f g. (forall a. Serialize a => Nonce24 -> Either (f a) (Assym (f a)) -> g a) -> OnionMessage f -> OnionMessage g
transcode f (OnionAnnounce a)        = OnionAnnounce        $ a { assymData = f (assymNonce a) (Right a) }
transcode f (OnionToRoute pub a)     = OnionToRoute pub     $ a { assymData = f (assymNonce a) (Right a) }
transcode f (OnionToRouteResponse a) = OnionToRouteResponse $ a { assymData = f (assymNonce a) (Right a) }
transcode f (OnionAnnounceResponse n8 n24 dta) = OnionAnnounceResponse n8 n24 $ f n24 $ Left dta