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

module Main where

import Types
import Hash
import Pty
import CmdLine
import Graphviz

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 qualified Data.ByteString.Lazy as L
import Data.List
import Data.List.NonEmpty (NonEmpty(..), toList)
import Data.Monoid
import Data.Aeson

main :: IO ()
main = do
	c <- getCmdLine
	case mode c of
		Test -> test
		Graphviz logfile -> graphviz logfile
 
test :: IO ()
test = do
	exitstatus <- go ">>> debug-me session starting"
	putStrLn ""
	putStrLn ">>> debug-me session is done"
	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

-- | 800 ms is about the latency to geosync orbit
networkDelay :: IO ()
networkDelay = threadDelay 800000 -- 150000 -- 800000

developer :: TChan (Activity Entered) -> TChan (Activity Seen) -> IO ()
developer ichan ochan = do
	startact <- atomically $ readTChan ochan
	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
		`concurrently` sendTtyOutput ochan devstate
	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 -> IO ()
sendTtyInput ichan devstate = go
  where
	go = do
		b <- B.hGetSome stdin 1024
		if b == B.empty
			then return ()
			else send b
	send b = do
		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'
		go

-- | Read activity from the TChan and display it to the developer.
sendTtyOutput :: TChan (Activity Seen) -> TVar DeveloperState -> IO ()
sendTtyOutput ochan devstate = go
  where
	go = do
		v <- atomically $ processOutput ochan devstate
		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
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
			Activity (Rejected _) _ _ -> do
 				writeTVar devstate ds'
				return Beep
		else return $ ProtocolError $ "Illegal Seen value: " ++ show (act, hash act) ++ "\n" ++ show ds

-- | 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
	let l = ActivitySeen (startact, hash startact)
	logger l
	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)

type Logger = ActivityLog -> IO ()

withLogger :: FilePath -> (Logger -> IO a) -> IO a
withLogger logfile a = withFile logfile WriteMode (a . mkLogger)

mkLogger :: Handle -> Logger
mkLogger h a = do
	L.hPut h (encode a)
	hPutStr h "\n"
	hFlush h

-- | 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
		act <- atomically $ do
			let seen = Seen (Val b)
			sendDeveloper ochan backlog (Proto seen)
		logger $ ActivitySeen (act, hash act)
		go

sendDeveloper :: TChan (Activity Seen) -> TVar Backlog -> Proto Seen -> STM (Activity Seen)
sendDeveloper ochan backlog pseen = do
	Backlog (bl@(prev :| _)) <- readTVar backlog
	let prevhash = activityLogHash prev
	let act = Activity pseen (Just prevhash) dummySignature
	writeTChan ochan act
	writeTVar backlog (Backlog (ActivitySeen (act, hash act) :| 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
		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) = truncateBacklog bl entered
			if isLegalEntered entered bl'
				then do
					let l = ActivityEntered (entered, hash entered)
					writeTVar backlog (Backlog (l :| toList bll))
					return (Right (entered, l))
				else do
					let reject = Rejected entered
					Left <$> sendDeveloper ochan backlog reject
		case v of
			Right (entered, l) -> do
				logger l
				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, hash rejact)
				go

-- | Truncate the Backlog to remove entries older than the one
-- that the Activity refers to.
--
-- 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 _)
	| Just (activityLogHash b) == hp = Backlog (b :| [])
	| otherwise = Backlog (b :| go [] l)
  where
	go c [] = reverse c
	go c (x:xs)
		| Just (activityLogHash x) == hp = reverse (x:c)
		| otherwise = go (x:c) 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 (activityLogHash lastseen) == hp = True
	| B.null (val (echoData entered)) = False -- optimisation
	| any (== hp) (map (Just . activityLogHash) bl) =
		let sincehp = reverse (lastseen : takeWhile (\l -> Just (activityLogHash l) /= hp) bl)
		in echoData entered == mconcat (map getseen 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 

dummySignature :: Signature
dummySignature = Signature mempty

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