{-# 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 (Message Seen) -> SessionID -> IO ()) -> UrlString -> IO ()
run' runner url = void $ runClientApp app
  where
	connect = ConnectMode (T.pack url)
	app = clientApp connect Developer userMessages runner

userMessages :: LogMessage -> Maybe (Message Seen)
userMessages (User m) = Just m
userMessages (Developer _) = Nothing

developer :: TMChan (Message Entered) -> TMChan (Message Seen) -> SessionID -> IO ()
developer ichan ochan _ = withLogger "debug-me-developer.log" $ \logger -> do
	devstate <- processSessionStart ochan logger
	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
	, sentSince :: [B.ByteString]
	, enteredSince :: [Hash]
	, lastActivity :: Hash
	, developerSessionKey :: MySessionKey
	, developerSigVerifier :: SigVerifier
	}

-- | 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 (Message Seen) -> TVar DeveloperState -> Logger -> IO ()
sendTtyOutput ochan devstate logger = go
  where
	go = do
		v <- atomically $ getUserMessage ochan devstate
		case v of
			Nothing -> return ()
			Just (o, msg) -> do
				logger $ User msg
				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 (Message Seen) -> 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 "No response from server to our session key")
			<$> atomically (getUserMessage ochan devstate)
		logger $ User 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

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

-- | Get messages from user, check their signature, and make sure that they
-- are properly chained from past messages, before returning.
getUserMessage :: TMChan (Message Seen) -> TVar DeveloperState -> STM (Maybe (Output, Message Seen))
getUserMessage ochan devstate = do
	mmsg <- readTMChan ochan
	case mmsg of
		Nothing -> return Nothing
		Just msg -> do
			ds <- readTVar devstate
			-- Check signature before doing anything else.
			if verifySigned (developerSigVerifier ds) msg
				then do
					o <- process ds msg
					return (Just (o, msg))
				else getUserMessage ochan devstate
  where
	process 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))
	process 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
	process _ (ControlMessage (Control c@(SessionKey _) _)) =
		return (GotControl c)
	process _ (ControlMessage (Control c@(SessionKeyAccepted _) _)) =
		return (GotControl c)
	process _ (ControlMessage (Control c@(SessionKeyRejected _) _)) =
		return (GotControl c)

-- | 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?
	| hp == lastSeen 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
			}
	-- 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
				}
  where
	acth = hash act
	yes ds' = (True, ds')

-- | Start by reading the initial two messages from the user side,
-- their session key and the startup message.
processSessionStart :: TMChan (Message Seen) -> Logger -> IO (TVar DeveloperState)
processSessionStart ochan logger = do
	sessionmsg <- fromMaybe (error "Did not get session initialization message")
		<$> atomically (readTMChan ochan)
	logger $ User sessionmsg
	sigverifier <- case sessionmsg of
		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 $ User startmsg
	starthash <- case startmsg of
		ActivityMessage act@(Activity (Seen (Val b)) Nothing _)
			| verifySigned sigverifier act -> do
				B.hPut stdout b
				hFlush stdout	
				return (hash act)
		_ -> error $ "Unexpected startup message: " ++ show startmsg
	sk <- genMySessionKey
	newTVarIO $ DeveloperState
		{ lastSeen = starthash
		, sentSince = mempty
		, enteredSince = mempty
		, lastActivity = starthash
		, developerSessionKey = sk
		, developerSigVerifier = sigverifier
		}