{-# LANGUAGE OverloadedStrings #-}

-- | Extract Lua programs embedded in other files
module Oughta.Extract
  ( SourceMap
  , sourceMapFile
  , lookupSourceMap
  , LuaProgram
  , programText
  , sourceMap
  , addPrefix
  , plainLuaProgram
  , fromLines
  , fromLineComments
  ) where

import Data.Foldable qualified as Foldable
import Data.IntMap.Strict (IntMap)
import Data.IntMap.Strict qualified as IntMap
import Data.Maybe qualified as Maybe
import Data.Text (Text)
import Data.Text qualified as Text

-- | Map from Lua line numbers to source line numbers
data SourceMap
  = SourceMap
    { SourceMap -> Text
sourceMapFile :: !Text
    , SourceMap -> IntMap Int
sourceMapLines :: !(IntMap Int)
    }

empty :: FilePath -> SourceMap
empty :: FilePath -> SourceMap
empty FilePath
path = Text -> IntMap Int -> SourceMap
SourceMap (FilePath -> Text
Text.pack FilePath
path) IntMap Int
forall a. IntMap a
IntMap.empty

lookupSourceMap ::
  -- | File path
  Text ->
  -- | Lua line number
  Int ->
  SourceMap ->
  Int
lookupSourceMap :: Text -> Int -> SourceMap -> Int
lookupSourceMap Text
path Int
luaLine (SourceMap Text
f IntMap Int
m) =
  if Text
path Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
f
  then Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
Maybe.fromMaybe Int
luaLine (Int -> IntMap Int -> Maybe Int
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
luaLine IntMap Int
m)
  else Int
luaLine

-- | Lua program paired with a 'SourceMap'
data LuaProgram
  = LuaProgram
    { LuaProgram -> Text
programText :: !Text
    , LuaProgram -> SourceMap
sourceMap :: !SourceMap
    }

-- | Add a prefix (i.e., prelude).
--
-- Bumps all subsequent line numbers
addPrefix :: Text -> LuaProgram -> LuaProgram
addPrefix :: Text -> LuaProgram -> LuaProgram
addPrefix Text
pfx LuaProgram
prog =
  let pfx' :: Text
pfx' = Text
pfx Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n" in
  LuaProgram
prog
  { programText = pfx' <> programText prog
  , sourceMap =
      let sm = LuaProgram -> SourceMap
sourceMap LuaProgram
prog in
      sm { sourceMapLines = IntMap.mapKeys (+ Text.count "\n" pfx') (sourceMapLines sm) }
  }

-- | A standalone Lua program
plainLuaProgram :: FilePath -> Text -> LuaProgram
plainLuaProgram :: FilePath -> Text -> LuaProgram
plainLuaProgram FilePath
path Text
txt = Text -> SourceMap -> LuaProgram
LuaProgram Text
txt (FilePath -> SourceMap
empty FilePath
path)

-- | Extract a Lua program embedded in certain lines of another file
fromLines :: FilePath -> (Text -> Maybe Text) -> Text -> LuaProgram
fromLines :: FilePath -> (Text -> Maybe Text) -> Text -> LuaProgram
fromLines FilePath
path Text -> Maybe Text
filt Text
txt =
  -- luaLineNo starts at 2 becaus we add a newline before each line
  let (Int
_, Text
t, SourceMap
m) = ((Int, Text, SourceMap) -> (Int, Text) -> (Int, Text, SourceMap))
-> (Int, Text, SourceMap)
-> [(Int, Text)]
-> (Int, Text, SourceMap)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' (Int, Text, SourceMap) -> (Int, Text) -> (Int, Text, SourceMap)
go (Int
2, Text
"", FilePath -> SourceMap
empty FilePath
path) ([Int] -> [Text] -> [(Int, Text)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1..] (Text -> [Text]
Text.lines Text
txt)) in
  Text -> SourceMap -> LuaProgram
LuaProgram Text
t SourceMap
m
  where
    go :: (Int, Text, SourceMap) -> (Int, Text) -> (Int, Text, SourceMap)
go (Int
luaLineNo, Text
t, SourceMap
sm) (Int
sourceLineNo, Text
line) =
      case Text -> Maybe Text
filt Text
line of
        Just Text
line' ->
          -- trace ("LINE '" ++ Text.unpack line' ++ " at " ++ show luaLineNo ++ " is " ++ show sourceLineNo) $
          ( Int
luaLineNo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
          , Text
t Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
line'
          , SourceMap
sm { sourceMapLines = IntMap.insert luaLineNo sourceLineNo (sourceMapLines sm) }
          )
        Maybe Text
Nothing ->
          ( Int
luaLineNo
          , Text
t
          , SourceMap
sm
          )

-- | Extract a Lua program embedded in the line comments of another file
fromLineComments ::
  FilePath ->
  -- | Start of comment marker, e.g., @"# "@
  Text ->
  -- | Whole file
  Text ->
  LuaProgram
fromLineComments :: FilePath -> Text -> Text -> LuaProgram
fromLineComments FilePath
path Text
c = FilePath -> (Text -> Maybe Text) -> Text -> LuaProgram
fromLines FilePath
path (Text -> Text -> Maybe Text
Text.stripPrefix Text
c)