{-# LANGUAGE OverloadedStrings, RankNTypes, FlexibleContexts #-}

module Main where

import Types
import Hash
import Pty
import Memory
import CmdLine
import Log
import Graphviz
import Replay
import Session
import Crypto

import Control.Concurrent
import Control.Concurrent.Async
import Control.Concurrent.STM
import System.IO
import System.Process
import System.Exit
import qualified Data.ByteString as B
import Data.List
import Data.List.NonEmpty (NonEmpty(..), toList)
import Data.Monoid
import Data.Time.Clock.POSIX

main :: IO ()
main = do
	c <- getCmdLine
	case mode c of
		Test -> test
		Graphviz g -> graphviz g
		Replay r -> replay r
 
test :: IO ()
test = do
	exitstatus <- go startSession
	sessionDone
	exitWith exitstatus
  where
	go startmsg = runWithPty "dash" [] $ \(p, ph) -> do
		ichan <- newTChanIO
		ochan <- newTChanIO
		dthread <- async (developer ichan ochan)
		uthread <- async (user startmsg p ichan ochan)
		exitstatus <- waitForProcess ph
		cancel dthread
		cancel uthread
		return exitstatus

networkDelay :: IO ()
networkDelay = threadDelay 800000 -- 800 ms ; the latency to geosync orbit
-- networkDelay = threadDelay 150000 -- 150 ms ; transatlantic latency

developer :: TChan (Activity Entered) -> TChan (Activity Seen) -> IO ()
developer ichan ochan = withLogger "debug-me-developer.log" $ \logger -> do
	startact <- atomically $ readTChan ochan
	logger $ ActivitySeen startact
	case startact of
		Activity (Proto (Seen (Val b))) Nothing sig -> do
			B.hPut stdout b
			hFlush stdout	
		_ -> protocolError $ "Unexpected startup: " ++ show startact
	devstate <- newTVarIO $ DeveloperState
		{ lastSeen = hash startact
		, sentSince = mempty
		, enteredSince = mempty
		, lastActivity = hash startact
		}
	_ <- sendTtyInput ichan devstate logger
		`concurrently` sendTtyOutput ochan devstate logger
	return ()

data DeveloperState = DeveloperState
	{ lastSeen :: Hash
	, sentSince :: [B.ByteString]
	, enteredSince :: [Hash]
	, lastActivity :: Hash
	}
	deriving (Show)

-- | Read things typed by the developer, and forward them to the TChan.
sendTtyInput :: TChan (Activity Entered) -> TVar DeveloperState -> Logger -> IO ()
sendTtyInput ichan devstate logger = go
  where
	go = do
		b <- B.hGetSome stdin 1024
		if b == B.empty
			then return ()
			else send b
	send b = do
		act <- atomically $ do
			ds <- readTVar devstate
			let ed = if lastActivity ds == lastSeen ds
				then B.concat $ sentSince ds
				else case reverse (sentSince ds) of
					[] -> mempty
					(lb:_) -> lb
			let entered = Entered
				{ enteredData = Val b
				, echoData = Val ed
				}
			let act = Activity (Proto entered) (Just $ lastActivity ds) dummySignature
			writeTChan ichan act
			let acth = hash act
			let ds' = ds
				{ sentSince = sentSince ds ++ [b]
				, enteredSince = enteredSince ds ++ [acth]
				, lastActivity = acth
				}
			writeTVar devstate ds'
			return act
		logger $ ActivityEntered act
		go

-- | Read activity from the TChan and display it to the developer.
sendTtyOutput :: TChan (Activity Seen) -> TVar DeveloperState -> Logger -> IO ()
sendTtyOutput ochan devstate logger = go
  where
	go = do
		(v, act) <- atomically $ processOutput ochan devstate
		logger $ ActivitySeen act
		case v of
			ProtocolError e -> protocolError e
			TtyOutput b -> do
				B.hPut stdout b
				hFlush stdout
				go
			Beep -> do
				B.hPut stdout "\a"
				hFlush stdout
				go

data Output = TtyOutput B.ByteString | Beep | ProtocolError String

processOutput :: TChan (Activity Seen) -> TVar DeveloperState -> STM (Output, Activity Seen)
processOutput ochan devstate = do
	act <- readTChan ochan
	ds <- readTVar devstate
	let (legal, ds') = isLegalSeen act ds
	if legal
		then case act of
			Activity (Proto (Seen (Val b))) _ _ -> do
				writeTVar devstate ds'
				return (TtyOutput b, act)
			Activity (Rejected _) _ _ -> do
 				writeTVar devstate ds'
				return (Beep, act)
		else return (ProtocolError $ "Illegal Seen value: " ++ show (act, hash act) ++ "\n" ++ show ds, act)

-- | Check if the Seen activity is legal, and returns an updated
-- DeveloperState.
isLegalSeen :: Activity Seen -> DeveloperState -> (Bool, DeveloperState)
isLegalSeen act@(Activity p (Just hp) sig) ds
	-- Does it chain to the last Seen value?
	| hp == lastSeen ds = case p of
		Rejected _ -> yesrej
		Proto (Seen (Val b)) ->
			-- Trim sentSince and enteredSince to
			-- values after the Seen value.
			let ss = sentSince ds
			    es = enteredSince ds
			    n = B.length b
			    (ss', es') = if b `B.isPrefixOf` mconcat ss
			    	then (drop n ss, drop n es)
				else (mempty, mempty)
			in yes $ DeveloperState
				{ lastSeen = acth
				, sentSince = ss'
				, enteredSince = es'
				, lastActivity = acth
				}
	-- Does it chain to something we've entered since the last Seen
	-- value? Eg, user sent A, we replied B C, and the user has
	-- now replied to B.
	-- If so, we can drop B (and anything before it) from
	-- enteredSince and sentSince.
	| otherwise = case elemIndex hp (enteredSince ds) of
		Nothing -> (False, ds)
		Just i -> case p of
			Rejected _ -> yesrej
			Proto (Seen (Val _)) ->
				let ss = sentSince ds
				    es = enteredSince ds
				    ss' = drop (i+1) ss
				    es' = drop (i+1) es
				in yes $ DeveloperState
					{ lastSeen = acth
					, sentSince = ss'
					, enteredSince = es'
					, lastActivity = acth
					}
  where
	acth = hash act
	yes ds' = (True, ds')
	-- When they rejected a message we sent, anything we sent
	-- subsequently will also be rejected, so forget about it.
	yesrej = yes $ ds
		{ lastSeen = acth
		, lastActivity = acth
		}
isLegalSeen (Activity _ Nothing _) ds = (False, ds)

user :: B.ByteString -> Pty -> TChan (Activity Entered) -> TChan (Activity Seen) -> IO ()
user startmsg p ichan ochan = withLogger "debug-me.log" $ \logger -> do
	let startact = Activity (Proto (Seen (Val (startmsg <> "\r\n")))) Nothing dummySignature
	logger $ ActivitySeen startact
	l <- mkActivityLog (ActivitySeen startact) <$> getPOSIXTime
	atomically $ writeTChan ochan startact
	backlog <- newTVarIO $ Backlog (l :| [])
	_ <- sendPtyOutput p ochan backlog logger
		`concurrently` sendPtyInput ichan ochan p backlog logger
	return ()

-- | Log of recent Activity, with the most recent first.
data Backlog = Backlog (NonEmpty ActivityLog)
	deriving (Show)

-- | Forward things written to the Pty out the TChan.
sendPtyOutput :: Pty -> TChan (Activity Seen) -> TVar Backlog -> Logger -> IO ()
sendPtyOutput p ochan backlog logger = go
  where
	go = do
		b <- readPty p
		now <- getPOSIXTime
		act <- atomically $ do
			let seen = Seen (Val b)
			sendDeveloper ochan backlog (Proto seen) now
		logger $ ActivitySeen act
		go

sendDeveloper :: TChan (Activity Seen) -> TVar Backlog -> Proto Seen -> POSIXTime -> STM (Activity Seen)
sendDeveloper ochan backlog pseen now = do
	Backlog (bl@(prev :| _)) <- readTVar backlog
	let prevhash = loggedHash prev
	let act = Activity pseen (Just prevhash) dummySignature
	let l = mkActivityLog (ActivitySeen act) now
	writeTChan ochan act
	writeTVar backlog $ Backlog (l :| toList bl)
	return act

-- | Read things to be entered from the TChan, verify if they're legal,
-- and send them to the Pty.
sendPtyInput :: TChan (Activity Entered) -> TChan (Activity Seen) -> Pty -> TVar Backlog -> Logger -> IO ()
sendPtyInput ichan ochan p backlog logger = go
  where
	go = do
		networkDelay
		now <- getPOSIXTime
		v <- atomically $ do
			entered <- readTChan ichan
			bl <- readTVar backlog
			-- Don't need to retain backlog before the Activity
			-- that entered references.
			let bl'@(Backlog bll) = reduceBacklog $
				truncateBacklog bl entered
			if isLegalEntered entered bl'
				then do
					let l = mkActivityLog (ActivityEntered entered) now
					writeTVar backlog (Backlog (l :| toList bll))
					return (Right entered)
				else do
					let reject = Rejected entered
					Left <$> sendDeveloper ochan backlog reject now
		case v of
			Right entered -> do
				logger (ActivityEntered entered)
				case activity entered of
					Proto e -> writePty p (val (enteredData e))
					Rejected r -> protocolError $ "User side received a Rejected: " ++ show r
				go
			Left rejact -> do
				logger $ ActivitySeen rejact
				go

-- | Truncate the Backlog to remove entries older than the one
-- that the Activity Entered refers to, but only if the referred
-- to Activity is an Activity Seen. 
--
-- Once the developer has referred to a given Activity Seen in
-- their Activity Entered, they cannot refer backwards to anything
-- that came before it.
--
-- If the Activity refers to an item not in the backlog, no truncation is
-- done.
truncateBacklog :: Backlog -> Activity Entered -> Backlog
truncateBacklog (Backlog (b :| l)) (Activity _ hp _)
	| truncationpoint b = Backlog (b :| [])
	| otherwise = Backlog (b :| go [] l)
  where
	go c [] = reverse c
	go c (x:xs)
		| truncationpoint x = reverse (x:c)
		| otherwise = go (x:c) xs
	truncationpoint x@(ActivityLog { loggedActivity = ActivitySeen {}}) = Just (loggedHash x) == hp
	truncationpoint _ = False

-- | To avoid DOS attacks that try to fill up the backlog and so use all
-- memory, don't let the backlog contain more than 1000 items, or
-- more than 16 megabytes of total data. (Excluding the most recent
-- item).
reduceBacklog :: Backlog -> Backlog
reduceBacklog (Backlog (b :| l)) = Backlog (b :| go 0 (take 1000 l))
  where
	go _ [] = []
	go n (x:xs)
		| n > 16777216 = []
		| otherwise = x : go (n + dataSize x) xs

-- | Entered activity is legal when it points to the last Seen activvity,
-- because this guarantees that the person who entered it saw
-- the current state of the system before manipulating it.
--
-- To support typeahead on slow links, some echoData may be provided
-- in the Entered activity. If the Entered activity points
-- to an older Seen activity, then the echoData must match the
-- concatenation of all Seen activities after that one, up to the
-- most recent Seen activity.
isLegalEntered :: Activity Entered -> Backlog -> Bool
isLegalEntered (Activity (Proto entered) hp sig) (Backlog (lastseen :| bl))
	| Just (loggedHash lastseen) == hp = True
	| B.null (val (echoData entered)) = False -- optimisation
	| any (== hp) (map (Just . loggedHash) bl) =
		let sincehp = reverse (lastseen : takeWhile (\l -> Just (loggedHash l) /= hp) bl)
		in echoData entered == mconcat (map (getseen . loggedActivity) sincehp)
	| otherwise = False
  where
	getseen (ActivitySeen a)  = case activity a of
		Proto s -> seenData s
		_ -> mempty
	getseen (ActivityEntered _) = mempty
-- Developer should never send Rejected.
isLegalEntered (Activity (Rejected _) _ _) _ = False 

-- | Temporary hack while user and developer share a process.
protocolError :: String -> IO a
protocolError e = do
	hPutStrLn stderr e
	hFlush stderr
	_ <- exitWith (ExitFailure 101)
	error e