{-
 - Copyright (C) 2015-2016 Ramakrishnan Muthukrishnan <ram@rkrishnan.org>
 -
 - This file is part of FuncTorrent.
 -
 - FuncTorrent is free software; you can redistribute it and/or modify
 - it under the terms of the GNU General Public License as published by
 - the Free Software Foundation; either version 3 of the License, or
 - (at your option) any later version.
 -
 - FuncTorrent is distributed in the hope that it will be useful,
 - but WITHOUT ANY WARRANTY; without even the implied warranty of
 - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 - GNU General Public License for more details.
 -
 - You should have received a copy of the GNU General Public License
 - along with FuncTorrent; if not,  see <http://www.gnu.org/licenses/>
 -}

{-# LANGUAGE OverloadedStrings #-}
module FuncTorrent.Tracker.Udp
       (trackerLoop
       ) where

import Control.Applicative (liftA2)
import Control.Monad (forever, void)
import Control.Concurrent (threadDelay)
import Control.Concurrent.MVar (readMVar, putMVar, isEmptyMVar, swapMVar)
import Control.Monad.Reader (ReaderT, runReaderT, ask, liftIO)
import Data.Binary (Binary(..), encode, decode)
import Data.Binary.Get (Get, isEmpty, getWord32be, getWord64be, getByteString)
import Data.Binary.Put (putWord16be, putWord64be, putWord32be, putByteString)
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as BC
import Data.ByteString.Lazy (fromStrict, toStrict)
import Data.Word (Word16, Word32, Word64)
import Network (PortNumber)
import Network.Socket (Socket, Family( AF_INET ), SocketType( Datagram ), defaultProtocol, SockAddr(..), socket, close, getAddrInfo, addrAddress, SockAddr(..))
import Network.Socket.ByteString (sendTo, recvFrom)
import System.Random (randomIO)
import System.Timeout (timeout)

import FuncTorrent.PeerMsgs (Peer(..))
import FuncTorrent.Tracker.Types (TrackerEventState(..), TState(..), UdpTrackerResponse(..))
import FuncTorrent.Utils (Port, toIP, toPort, getHostname, getPort)
import qualified FuncTorrent.FileSystem as FS (MsgChannel, Stats(..), getStats)

-- UDP tracker: http://bittorrent.org/beps/bep_0015.html
data Action = Connect
            | Announce
            | Scrape
            deriving (Show, Eq)

data UDPRequest = ConnectReq Word32
                | AnnounceReq Word64 Word32 ByteString String Word64 Word64 Word64 TrackerEventState Word16
                | ScrapeReq Integer Integer ByteString
                deriving (Show, Eq)

data UDPResponse = ConnectResp Word32 Word64 -- transaction_id connection_id
                 | AnnounceResp Word32 Word32 Word32 Word32 [Peer] -- transaction_id interval leechers seeders [(ip, port)]
                 | ScrapeResp Integer Integer Integer Integer
                 | ErrorResp Integer String
                 deriving (Show, Eq)

data UDPTrackerHandle = UDPTrackerHandle { sock :: Socket
                                         , addr :: SockAddr
                                         }

actionToInteger :: Action -> Integer
actionToInteger Connect  = 0
actionToInteger Announce = 1
actionToInteger Scrape   = 2

eventToInteger :: TrackerEventState -> Integer
eventToInteger None = 0
eventToInteger Completed = 1
eventToInteger Started = 2
eventToInteger Stopped = 3

instance Binary UDPRequest where
  put (ConnectReq transId) = do
    putWord64be 0x41727101980
    putWord32be $ fromIntegral (actionToInteger Connect)
    putWord32be (fromIntegral transId)
  put (AnnounceReq connId transId infohash peerId down left' up event port) = do
    putWord64be $ fromIntegral connId
    putWord32be $ fromIntegral (actionToInteger Announce)
    putWord32be $ fromIntegral transId
    putByteString infohash
    putByteString (BC.pack peerId)
    putWord64be (fromIntegral down)
    putWord64be (fromIntegral left')
    putWord64be (fromIntegral up)
    putWord32be $ fromIntegral (eventToInteger event)
    putWord32be 0
    putWord32be 0
    putWord32be 10
    putWord16be $ fromIntegral port
  put ScrapeReq {} = undefined
  get = undefined

instance Binary UDPResponse where
  put = undefined
  get = do
    a <- getWord32be -- action
    case a of
      0 -> liftA2 ConnectResp (fromIntegral <$> getWord32be) (fromIntegral <$> getWord64be)
      1 -> do
        tid' <- fromIntegral <$> getWord32be
        interval' <- fromIntegral <$> getWord32be
        l <- getWord32be -- leechers
        s <- getWord32be -- seeders
        ipportpairs <- getIPPortPairs -- [(ip, port)]
        return $ AnnounceResp tid' interval' l s ipportpairs
      2 -> do
        tid' <- fromIntegral <$> getWord32be
        _ <- getWord32be
        _ <- getWord32be
        _ <- getWord32be
        return $ ScrapeResp tid' 0 0 0
      3 -> do -- error response
        tid' <- fromIntegral <$> getWord32be
        bs  <- getByteString 4
        return $ ErrorResp tid' $ BC.unpack bs
      _ -> error ("unknown response action type: " ++ show a)

sendRequest :: UDPTrackerHandle -> ByteString -> IO ()
sendRequest h req = do
  n <- sendTo (sock h) req (addr h)
  -- sanity check with n?
  return ()

recvResponse :: UDPTrackerHandle -> IO UDPResponse
recvResponse h = do
  (bs, saddr) <- recvFrom (sock h) (16*1024)
  return $ decode $ fromStrict bs

connectRequest :: ReaderT UDPTrackerHandle IO Word32
connectRequest = do
  h <- ask
  tidi <- liftIO randomIO
  let pkt = encode $ ConnectReq tidi
  liftIO $ sendRequest h (toStrict pkt)
  return tidi

connectResponse :: Word32 -> ReaderT UDPTrackerHandle IO Word64
connectResponse tid = do
  h <- ask
  resp <- liftIO $ recvResponse h
  -- check if nbytes is at least 16 bytes long
  case resp of
    (ConnectResp tidr cid) ->
      if tidr == tid
      then do
        liftIO $ putStrLn "connect succeeded"
        return cid
      else
        return 0
    _                      -> return 0

announceRequest :: Word64 -> ByteString -> String -> Word64 -> Word64 -> Word64 -> Word16 -> ReaderT UDPTrackerHandle IO Word32
announceRequest cid infohash peerId up down left' port = do
  h <- ask
  tidi <- liftIO randomIO
  let pkt = encode $ AnnounceReq cid tidi infohash peerId down left' up None port
  liftIO $ sendRequest h (toStrict pkt)
  return tidi

announceResponse :: Word32 -> ReaderT UDPTrackerHandle IO UdpTrackerResponse
announceResponse tid = do
  h <- ask
  resp <- liftIO $ recvResponse h
  case resp of
    (AnnounceResp tidr interval' ss ls xs) ->
      if tidr == tid
      then do
        liftIO $ putStrLn "announce succeeded"
        return $ UdpTrackerResponse ls ss interval' xs
      else
        return $ UdpTrackerResponse 0 0 0 []
    _ -> return $ UdpTrackerResponse 0 0 0 []

getIPPortPairs :: Get [Peer]
getIPPortPairs = do
  empty <- isEmpty
  if empty
    then return []
    else do
    ip <- toIP <$> getByteString 4
    port <- toPort <$> getByteString 2
    ipportpairs <- getIPPortPairs
    return $ Peer ip port : ipportpairs

startSession :: String -> Port -> IO UDPTrackerHandle
startSession host port = do
  s <- socket AF_INET Datagram defaultProtocol
  addrinfos <- getAddrInfo Nothing (Just host) (Just (show port))
  let (SockAddrInet p ip) = addrAddress $ head addrinfos
  putStrLn "connected to tracker"
  return UDPTrackerHandle { sock = s
                          , addr = SockAddrInet (fromIntegral port) ip }
  
closeSession :: UDPTrackerHandle -> IO ()
closeSession (UDPTrackerHandle s _) = close s

trackerLoop :: String -> PortNumber -> String -> ByteString -> FS.MsgChannel -> TState -> IO ()
trackerLoop url sport peerId infohash fschan tstate = forever $ do
  st <- readMVar <$> FS.getStats fschan
  up <- fmap FS.bytesRead st
  down <- fmap FS.bytesWritten st
  handle <- startSession host port
  stats <- timeout (15*oneSec) $ worker handle up down
  case stats of
    Nothing -> closeSession handle
    Just stats' -> do
      ps <- isEmptyMVar $ connectedPeers tstate
      if ps
        then
        putMVar (connectedPeers tstate) (peers stats')
        else
        void $ swapMVar (connectedPeers tstate) (peers stats')
      threadDelay $ fromIntegral (interval stats') * oneSec
      return ()
  where
    oneSec = 1000000
    port = getPort url
    host = getHostname url
    worker handle up down = flip runReaderT handle $ do
      t1 <- connectRequest
      cid <- connectResponse t1
      t2 <- announceRequest cid infohash peerId (fromIntegral up) (fromIntegral down) (fromIntegral (left tstate)) (fromIntegral sport)
      announceResponse t2
