{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralisedNewtypeDeriving #-}

module Language.Fiddle.Compiler.ImportResolution
  ( resolveImports,
    getImportResolutionState,
    importResolutionPhase,
  )
where

import qualified Codec.Compression.GZip as GZip
import Control.Arrow (Arrow (second))
import Control.Monad (filterM, when)
import Control.Monad.State (put)
import Control.Monad.Trans.Maybe (MaybeT (MaybeT, runMaybeT))
import Control.Monad.Writer.Lazy (MonadTrans (lift), MonadWriter (tell), WriterT (..), execWriterT)
import Data.Aeson (eitherDecode, encode)
import qualified Data.ByteString.Lazy as BL
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe, isNothing)
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Typeable
import Language.Fiddle.Ast
import Language.Fiddle.Compiler
import Language.Fiddle.Compiler.Expansion ()
import Language.Fiddle.Internal.UnitInterface
import Language.Fiddle.Types
import Options.Applicative
import System.Directory
import System.FilePath
import Text.Printf (printf)

data Flags = Flags
  { importDirectories :: [FilePath],
    interfaceDirectory :: Maybe FilePath
  }

parseFlags :: Parser Flags
parseFlags =
  Flags
    <$> many
      ( strOption
          ( long "import"
              <> short 'I'
              <> metavar "DIRECTORY"
              <> help "Directory to add to the import search path"
          )
      )
    <*> optional
      ( strOption
          ( long "intf-dir"
              <> short 'i'
              <> metavar "INTF_DIR"
              <> help "Directory to dump interface files to."
          )
      )

importResolutionPhase ::
  ( FilePath ->
    IO ([Diagnostic], [Artifact], Maybe (TreeType FiddleUnit Parsed))
  ) ->
  ( TreeType FiddleUnit Parsed ->
    IO ([Diagnostic], [Artifact], Maybe (TreeType FiddleUnit Checked))
  ) ->
  CompilationPhase CurrentStage ImportsResolved
importResolutionPhase parseFile compileToChecked =
  CompilationPhase parseFlags (getImportResolutionState parseFile compileToChecked) resolveImports

type GlobalState = Bool

type LocalState = ResolvedImports

type M = Compile GlobalState

type Annot = Commented SourceSpan

newtype ResolvedImports = ResolvedImports
  { importMap :: Map Text ([Diagnostic], Maybe UnitInterface)
  }
  deriving newtype (Semigroup, Monoid)

type CurrentStage = Parsed

type I = Identity

instance CompilationStage CurrentStage where
  type StageAfter CurrentStage = ImportsResolved
  type StageMonad CurrentStage = M
  type StageState CurrentStage = LocalState
  type StageFunctor CurrentStage = Identity
  type StageAnnotation CurrentStage = Annot

resolveImports ::
  Flags ->
  ResolvedImports ->
  FiddleUnit CurrentStage I Annot ->
  Compile () (FiddleUnit ImportsResolved I Annot)
resolveImports _ localState t = do
  (b, a) <- subCompile False $ advanceStage localState t
  if b
    then compilationFailure
    else return a

-- | Mark the current compilation as failed, but allows the import resolution to
-- continue in order to allow all import failures to be reported at the same
-- time rather than as piecemeal.
markFatal :: Compile Bool ()
markFatal = put True

deriving instance AdvanceStage CurrentStage ObjTypeBody

deriving instance AdvanceStage CurrentStage DeferredRegisterBody

deriving instance AdvanceStage CurrentStage RegisterBody

deriving instance AdvanceStage CurrentStage AnonymousBitsType

deriving instance AdvanceStage CurrentStage BitType

deriving instance AdvanceStage CurrentStage (ConstExpression u)

deriving instance AdvanceStage CurrentStage EnumBody

deriving instance AdvanceStage CurrentStage EnumConstantDecl

deriving instance AdvanceStage CurrentStage RegisterBitsDecl

deriving instance AdvanceStage CurrentStage PackageBody

deriving instance AdvanceStage CurrentStage ObjTypeDecl

deriving instance AdvanceStage CurrentStage FiddleUnit

deriving instance AdvanceStage CurrentStage (Expression u)

deriving instance AdvanceStage CurrentStage RegisterBitsTypeRef

deriving instance AdvanceStage CurrentStage ObjType

deriving instance (AdvanceStage CurrentStage t) => AdvanceStage CurrentStage (Directed t)

deriving instance AdvanceStage CurrentStage FiddleDecl

instance AdvanceStage CurrentStage ImportStatement where
  advanceStage s (ImportStatement path list _ a) = do
    let what = Map.lookup path (importMap s)
        empty = UnitInterface mempty mempty

    v <- case what of
      Nothing -> do
        emitDiagnosticError "Failed to lookup imports (This is a bug)" a
        markFatal
        return empty
      Just (diags, val) -> do
        tell diags
        when (isNothing val) markFatal
        return $ fromMaybe empty val

    return $ ImportStatement path list (Present v) a

getImportResolutionState ::
  ( FilePath ->
    IO ([Diagnostic], [Artifact], Maybe (TreeType FiddleUnit Parsed))
  ) ->
  ( TreeType FiddleUnit Parsed ->
    IO ([Diagnostic], [Artifact], Maybe (TreeType FiddleUnit Checked))
  ) ->
  Flags ->
  FiddleUnit CurrentStage Identity Annot ->
  IO ([Diagnostic], Maybe ResolvedImports)
getImportResolutionState parseFile compileToChecked flags unit = do
  fmap
    ( lookForFailures
        . second Just
    )
    $ execWriterT
    $ walk doWalk unit ()
  where
    doWalk :: forall t'. (Walk t', Typeable t') => t' Identity Annot -> () -> WriterT ([Diagnostic], ResolvedImports) IO (WalkContinuation ())
    doWalk u () = do
      case () of
        ()
          | Just
              (ImportStatement {importPath = path, importStatementAnnot = (unCommented -> a)}) <-
              castTS u -> do
              (diagnostics, _, unitInterface) <-
                lift $
                  ioGetImportInterface a flags (Text.unpack path)

              tell
                ( [],
                  ResolvedImports $ Map.singleton path (diagnostics, unitInterface)
                )
        _ -> return ()
      return $ Continue ()

    castTS ::
      ( Typeable t',
        Typeable t,
        Typeable f,
        Typeable a
      ) =>
      t' f a ->
      Maybe (t CurrentStage f a)
    castTS = cast

    lookForFailures :: ([Diagnostic], Maybe a) -> ([Diagnostic], Maybe a)
    lookForFailures (diags, a) = do
      if any (\(Diagnostic e _ _) -> e == Error) diags
        then (diags, Nothing)
        else (diags, a)

    findInterfaceFile :: Maybe FilePath -> FilePath -> FilePath -> IO FilePath
    findInterfaceFile Nothing foundPath _ = return $ interfaceFile foundPath
    findInterfaceFile (Just dir) _ p = return $ interfaceFile $ dir </> p

    ioGetImportInterface :: SourceSpan -> Flags -> FilePath -> IO ([Diagnostic], [Artifact], Maybe UnitInterface)
    ioGetImportInterface srcSpan flags fp = runCompl $ do
      let imports = importDirectories flags
          intfDir = interfaceDirectory flags
      path <- findFileInImportPath srcSpan imports fp
      intf <- lift2 $ findInterfaceFile intfDir path fp
      -- let intf = interfaceFile path
      valid <- lift2 $ interfaceFileValid path intf

      let doFullCompile = do
            parsed <- bump (parseFile path)
            unitInterface <- addDependency path . unwrap . fiddleUnitInterface <$> bump (compileToChecked parsed)
            lift2 $ writeInterfaceFile intf unitInterface
            return unitInterface

      if valid
        then do
          e <- lift2 (readInterfaceFile intf)
          case e of
            Right val -> do
              needFullRecompile <- lift2 $ checkNeedFullRecompile intf val
              if needFullRecompile
                then doFullCompile
                else return val
            Left err -> do
              tell ([Diagnostic Warning "Doing recompile" srcSpan], [])
              tell ([Diagnostic Warning err srcSpan], [])
              doFullCompile
        else doFullCompile

    addDependency path unitInterface =
      unitInterface {dependencies = path : dependencies unitInterface}

    interfaceFile filePath = takeDirectory filePath </> takeBaseName filePath <.> "fdi"

    checkNeedFullRecompile intfFile (UnitInterface {dependencies = dependencies}) =
      anyM
        ( \depfile -> do
            timeDep <- getModificationTime depfile
            timeIntf <- getModificationTime intfFile
            return (timeIntf < timeDep)
        )
        dependencies

    readInterfaceFile intfile =
      eitherDecode . GZip.decompress <$> BL.readFile intfile

    writeInterfaceFile intfile val = do
      createDirectoryIfMissing True (takeDirectory intfile)
      BL.writeFile intfile $ GZip.compress (encode val)

    interfaceFileValid :: FilePath -> FilePath -> IO Bool
    interfaceFileValid originalPath intfPath = do
      exists <- doesFileExist intfPath
      if exists
        then do
          timeOrig <- getModificationTime originalPath
          timeIntf <- getModificationTime intfPath
          return (timeIntf > timeOrig)
        else return False

    findFileInImportPath :: SourceSpan -> [FilePath] -> FilePath -> Compl FilePath
    findFileInImportPath sourceSpan paths path = do
      canonicalPaths <- lift2 $ mapM (canonicalizePath . (++ ("/" ++ path))) paths
      realPaths <- lift2 $ filterM doesFileExist canonicalPaths

      case realPaths of
        [] -> do
          lift $ tell ([Diagnostic Error (printf "Cannot find %s on path" path) sourceSpan], [])
          MaybeT (return Nothing)
        (a : _) -> return a

bump :: IO ([Diagnostic], [Artifact], Maybe a) -> Compl a
bump x = do
  (diags, artifacts, ma) <- lift2 x
  lift $ tell (diags, artifacts)
  MaybeT (return ma)

lift2 :: (Monad m, MonadTrans t0, MonadTrans t1) => m a -> t0 (t1 m) a
lift2 = lift . lift

runCompl :: Compl a -> IO ([Diagnostic], [Artifact], Maybe a)
runCompl c = (\(x, (y, z)) -> (y, z, x)) <$> runWriterT (runMaybeT c)

type Compl a = MaybeT (WriterT ([Diagnostic], [Artifact]) IO) a

-- allM :: (Monad m) => (a -> m Bool) -> [a] -> m Bool
-- allM _ [] = return True
-- allM fn (a : as) = do
--   b <- fn a
--   if b then allM fn as else return False

anyM :: (Monad m) => (a -> m Bool) -> [a] -> m Bool
anyM _ [] = return True
anyM fn (a : as) = do
  b <- fn a
  if b then return True else anyM fn as