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

import Network.QueryResponse
import ToxCrypto hiding (encrypt,decrypt)
import qualified ToxCrypto
import DHTTransport (NodeInfo(..),NodeId(..),SendNodes,nodeInfo)

import Control.Arrow
import qualified Data.ByteString as B
         ;import Data.ByteString (ByteString)
import Data.Functor.Identity
import Data.Serialize            as S (Get, Put, Serialize, get, put, runGet)
import Data.Typeable
import Data.Word
import GHC.TypeLits
import Network.Socket

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

type UDPTransport = Transport String SockAddr ByteString


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

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))

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


onionToOwner assym ret3 saddr = do
    ni <- nodeInfo (NodeId $ senderKey assym) saddr
    return $ OnionToOwner ni ret3
-- data CookieAddress = WithoutCookie NodeInfo | CookieAddress Cookie SockAddr


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 (OnionAnnounceResponse <$> get <*> get <*> get) bs -- Announce Response
    0x86 -> response $ runGet (OnionToRouteResponse <$> getOnionAssym) bs        -- Onion Data Response
    _ -> right

encodeOnionAddr :: (OnionMessage Encrypted,OnionToOwner) -> (ByteString, SockAddr)
encodeOnionAddr = _todo

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 0) crypto saddr (forThem m)
            0x81 -> forward forThem bs $ handleOnionRequest (Proxy :: Proxy 1) crypto saddr (forThem m)
            0x82 -> forward forThem bs $ handleOnionRequest (Proxy :: Proxy 2) crypto saddr (forThem m)
            0x8c -> forward forThem bs $ handleOnionResponse (Proxy :: Proxy 3) crypto saddr (forThem m)
            0x8d -> forward forThem bs $ handleOnionResponse (Proxy :: Proxy 2) crypto saddr (forThem m)
            0x8e -> forward forThem bs $ handleOnionResponse (Proxy :: Proxy 1) crypto saddr (forThem m)
            _    -> 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

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

instance Serialize (OnionRequest n) where { get = _todo; put = _todo }
instance Serialize (OnionResponse n) where { get = _todo; put = _todo }

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

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

data Addressed a = Addressed { sockAddr :: SockAddr, unaddressed :: a }

data ReturnPath (n :: Nat) where
    NoReturnPath :: ReturnPath 0
    ReturnPath :: Nonce24 -> Encrypted (Addressed (ReturnPath n)) -> ReturnPath (n + 1)

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


-- 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 :: Nat) msg where
    NotForwarded :: msg -> Forwarding 0 msg
    Forwarding :: Assym (Encrypted (Addressed (Forwarding n msg))) -> Forwarding (n + 1) msg

handleOnionRequest :: KnownNat n => proxy n -> TransportCrypto -> SockAddr -> IO a -> OnionRequest n -> IO a
handleOnionRequest = _todo

handleOnionResponse :: KnownNat n => proxy n -> TransportCrypto -> SockAddr -> IO a -> OnionResponse n -> IO a
handleOnionResponse = _todo

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 S.Serialize AnnounceRequest where
    get = AnnounceRequest <$> S.get <*> S.get <*> S.get
    put (AnnounceRequest p s k) = S.put (p,s,k)

getOnionRequest :: Get (Assym (Encrypted msg), ReturnPath 3)
getOnionRequest = _todo

data KeyRecord = NotStored    Nonce32
               | SendBackKey  PublicKey
               | Acknowledged Nonce32

getPublicKey :: Get PublicKey
getPublicKey = _todo

putPublicKey :: PublicKey -> Put
putPublicKey = _todo

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 = VarSize $ \AnnounceResponse {} -> _todo

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
    , dataToRoute :: Encrypted (Word8,ByteString)
    }

instance Serialize DataToRoute where
    get = return $ DataToRoute _todo _todo
    put _ = return () -- todo

encrypt :: TransportCrypto -> OnionMessage Identity -> OnionToOwner -> (OnionMessage Encrypted, OnionToOwner)
encrypt crypto msg rpath = (transcode (encryptMessage crypto) msg, rpath)

encryptMessage :: Serialize a =>
                  TransportCrypto -> Nonce24 -> Either (Identity a) (Assym (Identity a)) -> Encrypted a
encryptMessage crypto n (Right a) = ToxCrypto.encrypt secret plain
 where
    secret = computeSharedSecret (transportSecret crypto) (senderKey a) n
    plain  = encodePlain $ runIdentity $ assymData a
encryptMessage crypto n (Left x) = ToxCrypto.encrypt secret plain
 where
    secret = computeSharedSecret (transportSecret crypto) _todo 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


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