{-# 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 System.IO
import qualified Data.ByteString as B
import qualified Data.Text as T
import Data.List

run :: DeveloperOpts -> IO ()
run os = runClientApp $ clientApp (ConnectMode (T.pack (debugUrl os))) developer

developer :: TChan (Message Entered) -> TChan (Message Seen) -> SessionID -> IO ()
developer ichan ochan _ = inRawMode $ withLogger "debug-me-developer.log" $ \logger -> do
	devstate <- processSessionStart ochan logger
	ok <- authUser ichan ochan devstate logger
	if ok
		then do
			_ <- sendTtyInput ichan devstate logger
				`concurrently` sendTtyOutput ochan devstate logger
			return ()
		else do
			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 TChan.
sendTtyInput :: TChan (Message Entered) -> TVar DeveloperState -> Logger -> IO ()
sendTtyInput ichan devstate logger = go
  where
	go = do
		print "in  sendTtyInput"
		b <- B.hGetSome stdin 1024
		print "sending from dev"
		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)
			writeTChan 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 TChan and display it to the developer.
sendTtyOutput :: TChan (Message Seen) -> TVar DeveloperState -> Logger -> IO ()
sendTtyOutput ochan devstate logger = go
  where
	go = do
		(o, msg) <- atomically $ getUserMessage ochan devstate
		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 :: TChan (Message Entered) -> TChan (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 $ writeTChan ichan msg
	logger $ Developer msg
	waitresp pk
  where
	waitresp pk = do
		(o, msg) <- 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 :: TChan (Message Seen) -> TVar DeveloperState -> STM (Output, Message Seen)
getUserMessage ochan devstate = do
	msg <- readTChan ochan
	ds <- readTVar devstate
	-- Check signature before doing anything else.
	if verifySigned (developerSigVerifier ds) msg
		then do
			o <- process ds msg
			return (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 :: TChan (Message Seen) -> Logger -> IO (TVar DeveloperState)
processSessionStart ochan logger = do
	sessionmsg <- atomically $ readTChan 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 <- atomically $ readTChan 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
		}