    {-# LANGUAGE DeriveDataTypeable #-}
    {-# LANGUAGE ViewPatterns #-}
    {-# LANGUAGE CPP #-}
    -- |
    -- Module      : Network.Socks5.Command
    -- License     : BSD-style
    -- Maintainer  : Vincent Hanquez <vincent@snarc.org>
    -- Stability   : experimental
    -- Portability : unknown
    --
    module Network.Socks5.Command
        ( establish
        , Connect(..)
        , Command(..)
        , connectIPV4
        , connectIPV6
        , connectDomainName
        -- * lowlevel interface
        , rpc
        , rpc_
        , sendSerialized
        , waitSerialized
        ) where

    import Control.Applicative
    import Control.Exception
    import Data.ByteString (ByteString)
    import qualified Data.ByteString as B
    import qualified Data.ByteString.Char8 as BC
    import Data.Serialize

    import Network.Socket (Socket, PortNumber, HostAddress, HostAddress6)
    import Network.Socket.ByteString

    import Network.Socks5.Types
    import Network.Socks5.Wire

    establish :: Socket -> [SocksMethod] -> IO SocksMethod
    establish socket methods = do
        sendAll socket (encode $ SocksHello methods)
        getSocksHelloResponseMethod <$> runGetDone get (recv socket 4096)

    newtype Connect = Connect SocksAddress deriving (Show,Eq,Ord)

    class Command a where
        toRequest   :: a -> SocksRequest
        fromRequest :: SocksRequest -> Maybe a

    instance Command SocksRequest where
        toRequest   = id
        fromRequest = Just

    instance Command Connect where
        toRequest (Connect (SocksAddress ha port)) = SocksRequest
                { requestCommand  = SocksCommandConnect
                , requestDstAddr  = ha
                , requestDstPort  = fromIntegral port
                }
        fromRequest req
            | requestCommand req /= SocksCommandConnect = Nothing
            | otherwise = Just $ Connect $ SocksAddress (requestDstAddr req) (requestDstPort req)

    connectIPV4 :: Socket -> HostAddress -> PortNumber -> IO (HostAddress, PortNumber)
    connectIPV4 socket hostaddr port = onReply <$> rpc_ socket (Connect $ SocksAddress (SocksAddrIPV4 hostaddr) port)
        where onReply (SocksAddrIPV4 h, p) = (h, p)
              onReply _                    = error "ipv4 requested, got something different"

    connectIPV6 :: Socket -> HostAddress6 -> PortNumber -> IO (HostAddress6, PortNumber)
    connectIPV6 socket hostaddr6 port = onReply <$> rpc_ socket (Connect $ SocksAddress (SocksAddrIPV6 hostaddr6) port)
        where onReply (SocksAddrIPV6 h, p) = (h, p)
              onReply _                    = error "ipv6 requested, got something different"

    -- TODO: FQDN should only be ascii, maybe putting a "fqdn" data type
    -- in front to make sure and make the BC.pack safe.
    connectDomainName :: Socket -> String -> PortNumber -> IO (SocksHostAddress, PortNumber)
    connectDomainName socket fqdn port = rpc_ socket $ Connect $ SocksAddress (SocksAddrDomainName $ BC.pack fqdn) port

    sendSerialized :: Serialize a => Socket -> a -> IO ()
    sendSerialized sock a = sendAll sock $ encode a

    waitSerialized :: Serialize a => Socket -> IO a
    waitSerialized sock = runGetDone get (getMore sock)

    rpc :: Command a => Socket -> a -> IO (Either SocksError (SocksHostAddress, PortNumber))
    rpc socket req = do
        sendSerialized socket (toRequest req)
        onReply <$> runGetDone get (getMore socket)
        where onReply res@(responseReply -> reply) =
                    case reply of
                        SocksReplySuccess -> Right (responseBindAddr res, fromIntegral $ responseBindPort res)
                        SocksReplyError e -> Left e

    rpc_ :: Command a => Socket -> a -> IO (SocksHostAddress, PortNumber)
    rpc_ socket req = rpc socket req >>= either throwIO return

    -- this function expect all the data to be consumed. this is fine for intertwined message,
    -- but might not be a good idea for multi messages from one party.
    runGetDone :: Serialize a => Get a -> IO ByteString -> IO a
    runGetDone getter ioget = ioget >>= return . runGetPartial getter >>= r where
    #if MIN_VERSION_cereal(0,4,0)
        r (Fail s _)     = error s
    #else
        r (Fail s)       = error s
    #endif
        r (Partial cont) = ioget >>= r . cont
        r (Done a b)
            | not $ B.null b = error "got too many bytes while receiving data"
            | otherwise      = return a

    getMore :: Socket -> IO ByteString
    getMore socket = recv socket 4096