module Sketch.FRP.Copilot.Internals where

import Sketch.FRP.Copilot.Types
import Language.Copilot
import Control.Monad.Writer
import Control.Monad.State.Strict
import Data.Functor.Identity
import qualified Data.Map as M
import qualified Data.Set as S
import Data.Maybe

getTriggerLimit :: TriggerLimit -> Behavior Bool
getTriggerLimit (TriggerLimit b) = b
getTriggerLimit NoTriggerLimit = true

addTriggerLimit :: TriggerLimit -> Behavior Bool -> Behavior Bool
addTriggerLimit tl c = getTriggerLimit (tl <> TriggerLimit c)

-- | Gets a unique id.
getUniqueId :: String -> GenSketch ctx UniqueId
getUniqueId s = do
	UniqueIds m <- get
	let u = maybe 1 succ (M.lookup s m)
	put $ UniqueIds $ M.insert s u m
	return (UniqueId u)

-- | Generates a unique name.
uniqueName :: String -> UniqueId -> String
uniqueName s (UniqueId i)
	| i Prelude.== 1 = s
	| otherwise = s <> "_" <>  show i

uniqueName' :: String -> UniqueId -> String
uniqueName' s (UniqueId i) = s <> "_" <>  show i

-- | Use to create an empty framework. 
--
-- It helps to specify the type of context to use:
--
-- > (emptyFramework @Arduino) { ... }
emptyFramework :: Context ctx => GenFramework ctx
emptyFramework = mempty

mkCChunk :: [CLine] -> [CChunk]
mkCChunk l = [CChunk l]

-- | Copilot only supports calling a trigger with a given name once
-- per Spec; the generated C code will fail to build if the same name is
-- used in two triggers. This generates a unique alias that can be 
-- used in a trigger.
defineTriggerAlias
	:: String
	-> GenFramework ctx
	-> GenSketch ctx (GenFramework ctx, String)
defineTriggerAlias = defineTriggerAlias' ""

defineTriggerAlias'
	:: String
	-> String
	-> GenFramework ctx
	-> GenSketch ctx (GenFramework ctx, String)
defineTriggerAlias' suffix cfuncname f = do
	let basetname = if null suffix 
		then cfuncname
		else cfuncname <> "_" <> suffix
	u <- getUniqueId basetname
	let triggername = uniqueName basetname u
	let define = if cfuncname Prelude./= triggername
		then mkCChunk [ CLine $ "#define " <> triggername <> " " <> cfuncname	]
		else mempty
	return (f { defines = define <> defines f }, triggername)

data MkInputSource ctx t = InputSource
	{ defineVar :: [CChunk]
	-- ^ Added to the `Framework`'s `defines`, this typically
	-- defines a C variable.
	, setupInput :: [CChunk]
	-- ^ How to set up the input, not including pin mode.
	, inputPinmode :: M.Map ctx PinMode
	-- ^ How pins are used by the input.
	, readInput :: [CChunk]
	-- ^ How to read a value from the input, this typically
	-- reads a value into a C variable.
	, inputStream :: Stream t
	-- ^ How to use Copilot's extern to access the input values.
	}

mkInput :: MkInputSource ctx t -> GenSketch ctx (Behavior t)
mkInput i = do
	u <- getUniqueId "input"
	tell [(mkspec u, f u)]
	return (inputStream i)
  where
	f u ratelimited = Framework
		{ defines = defineVar i <> mkdefine u ratelimited
		, setups = setupInput i
		, earlySetups = mempty
		, pinmodes = M.map S.singleton (inputPinmode i)
		, loops = mkloops u ratelimited (readInput i)
		}

	varname = uniqueName "update_input"
	triggername = uniqueName "input"
	
	mkdefine _ NoTriggerLimit = []
	mkdefine u (TriggerLimit _) = mkCChunk $ map CLine
		[ "bool " <> varname u <> " = true;"
		, "void " <> triggername u <> " (bool v) {"
		, "  " <> varname u <> " = v;"
		, "}"
		]
	
	mkloops _ NoTriggerLimit reader = reader
	mkloops u (TriggerLimit _) reader = mkCChunk $ concat
		[ [ CLine $ "if (" <> varname u <> ") {" ]
		, map (\(CLine l) -> CLine $ "  " <> l ) readerlines
		, [ CLine "}" ]
		]
	  where
		readerlines = concatMap (\(CChunk l) -> l) reader

	mkspec _ NoTriggerLimit = return ()
	mkspec u (TriggerLimit c) = trigger (triggername u) true [arg c]

evalSketch :: Context ctx => GenSketch ctx a -> (Maybe Spec, GenFramework ctx)
evalSketch (GenSketch s) = (spec, f)
  where
	(is, fs) = unzip $ 
		runIdentity $ evalStateT (execWriterT s) (UniqueIds mempty)
	f = mconcat (map (\f' -> f' NoTriggerLimit) fs)
	-- Copilot will throw an ugly error if given a spec that does
	-- nothing at all, so return Nothing to avoid that.
	spec :: Maybe Spec
	spec = if null is
		then Nothing
		else Just $ sequence_ $ map (\i -> i NoTriggerLimit) is

-- | Extracts a copilot `Spec` from a `Sketch`.
--
-- This can be useful to intergrate with other libraries 
-- such as copilot-theorem.
sketchSpec :: Context ctx => GenSketch ctx a -> Spec
sketchSpec = fromMaybe (return ()) . fst . evalSketch