{-# LANGUAGE BangPatterns #-}

module Copilot.Zephyr.Main (zephyr) where

import Language.Copilot (Spec, interpret, reify)
import Copilot.Compile.C99 (compile)
import Copilot.Zephyr.Internals
import System.IO
import System.Directory
import System.IO.Temp (withSystemTempDirectory, createTempDirectory)
import System.FilePath
import System.Exit
import Data.List (isInfixOf, intercalate)
import Data.Maybe
import qualified Data.Map as M
import qualified Data.Set as S
import Options.Applicative

-- | Typically your program's main will be implemented using this.
-- For example:
--
-- > {-# LANGUAGE RebindableSyntax #-}
-- > 
-- > import Copilot.Zephyr.Board.Generic
-- > 
-- > main = zephyr $ do
-- >	led0 =: flashing
-- > 	delay =: MilliSeconds (constant 100)
--
-- Running this program compiles the `Sketch` into C code using copilot
-- and generates a Zephyr app in the directory "generated". That app
-- can be built and uploaded to your board using Zephyr, the same as any
-- other Zephyr app. See Zephyr's documentation for details.
--
-- This also supports interpreting a `Sketch`, without loading it onto a
-- board. Run the program with parameters "-i 4" to display what it
-- would do on the first 4 iterations. The output will look something like
-- this:
--
-- > gpio_pin_set_led0:  k_msleep:
-- > (false)             (100)
-- > (true)              (100)
-- > (false)             (100)
-- > (true)              (100)
zephyr :: Sketch () -> IO ()
zephyr s = go =<< execParser opts
  where
	opts = info (parseCmdLine <**> helper)
		( fullDesc
		<> progDesc "Run this program with no options to generate an Zephyr program."
		)

	go o = case (mspec, interpretSteps o) of
		(Nothing, _) -> do
			hPutStrLn stderr "This Sketch does not do anything."
			exitFailure
		(Just spec, Just n) -> interpret n spec
		(Just spec, Nothing) -> do
			let name = "generated"
			writeMain name spec f'
			writePrjConf name kconfigs
			writeAppOverlay name devicetree
			writeCMakeLists name
	
	(mspec, f) = evalSketch s
	
	-- Strict evaluation because this may throw errors,
	-- and we want to throw them before starting compilation.
	!(f', kconfigs, devicetree) = finalizeFramework f

data CmdLine = CmdLine
	{ interpretSteps :: Maybe Integer
	}


parseCmdLine :: Parser CmdLine
parseCmdLine = CmdLine
	<$> optional (option auto
		( long "interpret"
		<> short 'i'
		<> help "use copilot to interpret the program, displaying what it would do"
		<> metavar "NUM"
		))

writeMain :: String -> Spec -> Framework -> IO ()
writeMain name spec framework = do
	-- This could be a lot prettier, unfortunately copilot only exports
	-- an interface that writes the generated code to a file.
	-- And, the .c file includes a .h file that will make it fail to
	-- build, so that include has to be filtered out.
	withSystemTempDirectory "copilot" $ \toptmpdir -> do
		mytmpdir <- createTempDirectory toptmpdir "copilot"
		reify spec >>= compile (mytmpdir </> "copilot")
		c <- lines <$> readFile (mytmpdir </> "copilot.c")
		let c' = filter (Prelude.not . isInfixOf "#include \"") c
		createDirectoryIfMissing False name
		writeFile (name </> "main.c") $ 
			unlines $ map fromCLine $
				sketchFramework framework (map CLine c')

writeCMakeLists :: String -> IO ()
writeCMakeLists name = do
	createDirectoryIfMissing False name
	writeFile (name </> "CMakeLists.txt") $ unlines
		[ "cmake_minimum_required(VERSION 3.20.0)"
		, "find_package(Zephyr REQUIRED HINTS $ENV{ZEPHYR_BASE})"
		, "project(" ++ name ++ ")"
		, ""
		, "target_sources(app PRIVATE main.c)"
		]

data KConfig = KConfig String

writePrjConf :: String -> [KConfig] -> IO ()
writePrjConf name kconfigs = do
	createDirectoryIfMissing False name
	writeFile (name </> "prj.conf") $ unlines $
		map (\(KConfig s) -> s) kconfigs

newtype DeviceTree = DeviceTree String

writeAppOverlay :: String -> DeviceTree -> IO ()
writeAppOverlay name (DeviceTree s) = do
	createDirectoryIfMissing False name
	writeFile (name </> "app.overlay") s

-- Makes an Zephyr C program, using a Framework, and a list of lines of C
-- code generated by Copilot.
sketchFramework :: Framework -> [CLine] -> [CLine]
sketchFramework f ccode = map CLine $ concat
	[
		[ "/* automatically generated, do not edit */"
		, blank
		, "#include <zephyr.h>"
		, blank
		]
	, map fromCLine (fromchunks (defines f))
	, [blank]
	, map fromCLine ccode
	, [blank]
	,
		[ "void main(void)"
		]
	, codeblock $ concat
		[ map fromCLine (fromchunks (earlySetups f))
		, map fromCLine (fromchunks (setups f))
		, ["while (1)"]
		, codeblock $ concat
			[ map fromCLine (fromchunks (loops f))
			, [ "step();" ]
			]
		]
	]
  where
	blank = ""
	indent l = "  " <> l
	codeblock l = ["{"] <> map indent l <> ["}"]
	-- If two CChunks are identical, only include it once.
	-- This can happen when eg, the same setup code is generated
	-- to use a resource that's accessed more than once in a program.
	-- Note: Does not preserve order of chunks in the list.
	fromchunks :: [CChunk] -> [CLine]
	fromchunks cl = concatMap (\(CChunk l) -> l) $
		S.toList $ S.fromList cl

finalizeFramework :: Framework -> (Framework, [KConfig], DeviceTree)
finalizeFramework f = 
	let (pindefines, pinsetups) = unzip $ mapMaybe setuppin
	    	(M.toList (pinmodes f))
	    includes = if not (null (concat pindefines))
		then mkCChunk 
			[ CLine "#include <device.h>"
			, CLine "#include <devicetree.h>"
			, CLine "#include <drivers/gpio.h>"
			]
		else []
	    f' = f
		{ defines = includes <> concat pindefines <> defines f
		, setups = concat pinsetups <> setups f
		}
	    kconfigs = if not (null (concat pindefines))
		then [KConfig "CONFIG_GPIO=y"]
		else []
	    devicetree = DeviceTree $ unlines $
	    	concatMap mkdevicetree (M.keys (pinmodes f))
	   in (f', kconfigs, devicetree)
  where
	setuppin (Zephyr (GPIOAlias n) _, s)
		| s == S.singleton OutputMode = Just $
			( definepin n
			, setmode n (ored ["GPIO_OUTPUT", gpioflags n])
			)
		| s == S.singleton InputMode = Just $
			( definepin n
			, setmode n (ored ["GPIO_INPUT", gpioflags n])
			)
		| s == S.singleton InputPullupMode = Just $
			( definepin n
			, setmode n (ored ["GPIO_PULL_UP", gpioflags n])
			)
		| s == S.fromList [InputMode, InputPullupMode] = Just $
			( definepin n
			, setmode n (ored ["GPIO_INPUT", "GPIO_PULL_UP", gpioflags n])
			)
		| S.null s = Nothing
		| otherwise = error $
			"The program uses pin " ++ show n ++
			" in multiple ways in different places (" ++
			unwords (map show (S.toList s)) ++ "). " ++
			"This is not currently supported by zephyr-copilot."
	
	definepin n = mkCChunk
		[ CLine $ "#define " <> gpionode n <> " DT_ALIAS(" <> n <> ")"
		, CLine $ "#if DT_NODE_HAS_STATUS(" <> gpionode n <> ", okay)"
		, CLine $ "#define " <> gpiolabel n 
			<> " DT_GPIO_LABEL(" <> gpionode n <> ", gpios)"
		, CLine $ "#define " <> gpiopin n 
			<> " DT_GPIO_PIN(" <> gpionode n <> ", gpios)"
		, CLine $ "#define " <> gpioflags n 
			<> " DT_GPIO_FLAGS(" <> gpionode n <> ", gpios)"
		, CLine $ "#else"
		, CLine $ "#error \"Unsupported board: devicetree alias is not" 
			<> " defined for " <> n <> "\""
		, CLine $ "#define " <> gpiolabel n <> ""
		, CLine $ "#define " <> gpiopin n <> " 0"
		, CLine $ "#define " <> gpioflags n <> " 0"
		, CLine $ "#endif"
		, CLine $ "const struct device *" <> pinDevVar n <> ";"
		]

	setmode n v = mkCChunk
		[ CLine $ pinDevVar n <> " = device_get_binding(" 
			<> gpiolabel n <> ");"
		, CLine $ "gpio_pin_configure(" 
			<> pinDevVar n <> ", "
			<> gpiopin n <> ", " 
			<> v <> ");" 
		]
	
	gpiopin n = pinDevDef n
	gpiolabel n = pinDevDef (n <> "_LABEL")
	gpioflags n = pinDevDef (n <> "_FLAGS")
	gpionode n = pinDevNode n
	
	ored = intercalate " | "
	
	mkdevicetree (Zephyr _ GPIOAddressBuiltIn) = []
	mkdevicetree (Zephyr (GPIOAlias n) (GPIOAddress addr)) =
		[ "/ {"
		, "\tmypins {"
		, "\t\tcompatible = \"gpio-keys\";"
		, "\t\t" ++ n ++ ": " ++ n ++ "{"
		, "\t\t\tgpios = <&" ++ addr ++ " 0>;"
		, "\t\t\tlabel = \"" ++ n ++ "\";"
		, "\t\t};"
		, "\t};"
		, "\taliases {"
		, "\t\t" ++ map toalias n ++ " = &" ++ n ++ ";"
		, "\t};"
		, "};"
		]
	  where
		toalias '_' = '-'
		toalias c = c