{-# LANGUAGE OverloadedStrings #-}

module Role.Developer where

import Types
import Hash
import Log
import Crypto
import CmdLine
import WebSockets
import SessionID
import Pty

import Control.Concurrent.Async
import Control.Concurrent.STM
import Control.Concurrent.STM.TMChan
import System.IO
import qualified Data.ByteString as B
import qualified Data.Text as T
import Data.List
import Data.Maybe
import Control.Monad

run :: DeveloperOpts -> IO ()
run = run' developer . debugUrl

run' :: (TMChan (Message Entered) -> TMChan LogMessage -> SessionID -> IO ()) -> UrlString -> IO ()
run' runner url = void $ runClientApp app
  where
	connect = ConnectMode (T.pack url)
	app = clientApp connect Developer Just runner

developer :: TMChan (Message Entered) -> TMChan LogMessage -> SessionID -> IO ()
developer ichan ochan _ = withLogger "debug-me-developer.log" $ \logger -> do
	(devstate, startoutput) <- processSessionStart ochan logger
	emitOutput startoutput
	ok <- authUser ichan ochan devstate logger
	if ok
		then inRawMode $ void $ 
			sendTtyInput ichan devstate logger
				`race` sendTtyOutput ochan devstate logger
		else hPutStrLn stderr "\nUser did not grant access to their terminal."

data DeveloperState = DeveloperState
	{ lastSeen :: Hash
	-- ^ Last Seen value received from the user.
	, sentSince :: [B.ByteString]
	-- ^ Keys pressed since last Seen.
	, enteredSince :: [Hash]
	-- ^ Messages we've sent since the last Seen.
	, lastActivity :: Hash
	-- ^ Last message sent or received.
	, fromOtherDevelopersSince :: [Hash]
	-- ^ Messages received from other developers since the last Seen.
	-- (The next Seen may chain from one of these.)
	, developerSessionKey :: MySessionKey
	-- ^ Our session key.
	, userSigVerifier :: SigVerifier
	-- ^ Used to verify signatures on messages from the user.
	, developerSigVerifier :: SigVerifier
	-- ^ Used to verify signatures on messages from other developers.
	}

-- | Read things typed by the developer, and forward them to the TMChan.
sendTtyInput :: TMChan (Message 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 = mkSigned (developerSessionKey ds) $
				Activity entered (Just $ lastActivity ds)
			writeTMChan ichan (ActivityMessage act)
			let acth = hash act
			let ds' = ds
				{ sentSince = sentSince ds ++ [b]
				, enteredSince = enteredSince ds ++ [acth]
				, lastActivity = acth
				}
			writeTVar devstate ds'
			return act
		logger $ Developer $ ActivityMessage act
		go

-- | Read activity from the TMChan and display it to the developer.
sendTtyOutput :: TMChan LogMessage -> TVar DeveloperState -> Logger -> IO ()
sendTtyOutput ochan devstate logger = go
  where
	go = do
		v <- atomically $ getServerMessage ochan devstate
		case v of
			Nothing -> return ()
			Just (o, l) -> do
				logger l
				emitOutput o
				go

-- | Present our session key to the user. 
-- Wait for them to accept or reject it, while displaying any Seen data
-- in the meantime.
authUser :: TMChan (Message Entered) -> TMChan LogMessage -> TVar DeveloperState -> Logger -> IO Bool
authUser ichan ochan devstate logger = do
	ds <- atomically $ readTVar devstate
	pk <- myPublicKey (developerSessionKey ds)
	let msg = ControlMessage $ mkSigned (developerSessionKey ds) 
		(Control (SessionKey pk))
	atomically $ writeTMChan ichan msg
	logger $ Developer msg
	waitresp pk
  where
	waitresp pk = do
		(o, msg) <- fromMaybe (error "Looks like that debug-me session is over.")
			<$> atomically (getServerMessage ochan devstate)
		logger msg
		emitOutput o
		case o of
			GotControl (SessionKeyAccepted pk')
				| pk' == pk -> return True
			GotControl (SessionKeyRejected pk')
				| pk' == pk -> return False
			_ -> waitresp pk

data Output
	= TtyOutput B.ByteString 
	| Beep
	| ProtocolError String
	| GotControl ControlAction
	| NoOutput

emitOutput :: Output -> IO ()
emitOutput (ProtocolError e) =
	error ("Protocol error: " ++ e)
emitOutput (TtyOutput b) = do
	B.hPut stdout b
	hFlush stdout
emitOutput Beep = do
	B.hPut stdout "\a"
	hFlush stdout
emitOutput (GotControl _) =
	return ()
emitOutput NoOutput =
	return ()

-- | Get messages from server, check their signature, and make sure that they
-- are properly chained from past messages, before returning.
getServerMessage :: TMChan LogMessage -> TVar DeveloperState -> STM (Maybe (Output, LogMessage))
getServerMessage ochan devstate = do
	let ignore = getServerMessage ochan devstate
	mmsg <- readTMChan ochan
	case mmsg of
		Nothing -> return Nothing
		Just (User msg) -> do
			ds <- readTVar devstate
			-- Check user's signature before doing anything else.
			if verifySigned (userSigVerifier ds) msg
				then do
					o <- processuser ds msg
					return (Just (o, User msg))
				else ignore
		-- When other developers connect, learn their SessionKeys.
		Just (Developer msg@(ControlMessage (Control (SessionKey pk) _))) -> do
			let sigverifier = mkSigVerifier pk
			if verifySigned sigverifier msg
				then do
					ds <- readTVar devstate
					let sv = developerSigVerifier ds
					let sv' = sv `mappend` sigverifier
					writeTVar devstate $ ds
						{ developerSigVerifier = sv'
						}
					processdeveloper ds msg
					return (Just (NoOutput, Developer msg))
				else ignore
		Just (Developer msg) -> do
			ds <- readTVar devstate
			if verifySigned (developerSigVerifier ds) msg
				then do
					processdeveloper ds msg
					return (Just (NoOutput, Developer msg))
				else ignore
  where
	processuser ds (ActivityMessage act@(Activity (Seen (Val b)) _ _)) = do
		let (legal, ds') = isLegalSeen act ds
		if legal
			then do
				writeTVar devstate ds'
				return (TtyOutput b)
			else return (ProtocolError $ "Illegal Seen value: " ++ show (act, hash act))
	processuser ds (ControlMessage (Control (Rejected _) _)) = do
		-- When they rejected a message we sent,
		-- anything we sent subsequently will
		-- also be rejected, so forget about it.
		let ds' = ds
			{ sentSince = mempty
			, enteredSince = mempty
			}
		writeTVar devstate ds'
		return Beep
	processuser _ (ControlMessage (Control c@(SessionKey _) _)) =
		return (GotControl c)
	processuser _ (ControlMessage (Control c@(SessionKeyAccepted _) _)) =
		return (GotControl c)
	processuser _ (ControlMessage (Control c@(SessionKeyRejected _) _)) =
		return (GotControl c)

	processdeveloper ds (ActivityMessage a) = do
		let msghash = hash a
		let ss = msghash : fromOtherDevelopersSince ds
		writeTVar devstate (ds { fromOtherDevelopersSince = ss })
	processdeveloper _ (ControlMessage _) = return ()

-- | Check if the Seen activity is legal, forming a chain with previous
-- ones, and returns an updated DeveloperState.
--
-- Does not check the signature.
isLegalSeen :: Activity Seen -> DeveloperState -> (Bool, DeveloperState)
isLegalSeen (Activity _ Nothing _) ds = (False, ds)
isLegalSeen act@(Activity (Seen (Val b)) (Just hp) _) ds
	-- Does it chain to the last Seen activity or to
	-- something sent by another developer since the last Seen?
	| hp == lastSeen ds || hp `elem` fromOtherDevelopersSince ds = 
		-- 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 ds
			{ lastSeen = acth
			, sentSince = ss'
			, enteredSince = es'
			, lastActivity = acth
			, fromOtherDevelopersSince = mempty
			}
	-- 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 -> 
			let ss = sentSince ds
			    es = enteredSince ds
			    ss' = drop (i+1) ss
			    es' = drop (i+1) es
			in yes ds
				{ lastSeen = acth
				, sentSince = ss'
				, enteredSince = es'
				, lastActivity = acth
				, fromOtherDevelopersSince = mempty
				}
  where
	acth = hash act
	yes ds' = (True, ds')

-- | Start by reading the initial two messages from the user,
-- their session key and the startup message.
processSessionStart :: TMChan LogMessage -> Logger -> IO (TVar DeveloperState, Output)
processSessionStart ochan logger = do
	sessionmsg <- fromMaybe (error "Did not get session initialization message")
		<$> atomically (readTMChan ochan)
	logger sessionmsg
	sigverifier <- case sessionmsg of
		User (ControlMessage c@(Control (SessionKey pk) _)) ->
			let sv = mkSigVerifier pk
			in if verifySigned sv c
				then return sv
				else error "Badly signed session initialization message"
		_ -> error $ "Unexpected session initialization message: " ++ show sessionmsg
	startmsg <- fromMaybe (error "Did not get session startup message")
		<$> atomically (readTMChan ochan)
	logger startmsg
	let (starthash, output) = case startmsg of
		User (ActivityMessage act@(Activity (Seen (Val b)) Nothing _))
			| verifySigned sigverifier act ->
				(hash act, TtyOutput b)
		_ -> error $ "Unexpected startup message: " ++ show startmsg
	sk <- genMySessionKey
	st <- newTVarIO $ DeveloperState
		{ lastSeen = starthash
		, sentSince = mempty
		, enteredSince = mempty
		, lastActivity = starthash
		, fromOtherDevelopersSince = mempty
		, developerSessionKey = sk
		, userSigVerifier = sigverifier
		, developerSigVerifier = mempty
		}
	return (st, output)