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

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

import Control.Monad.Identity (Identity)
import Control.Monad.Writer.Lazy (MonadTrans (lift), MonadWriter (tell), WriterT (WriterT), execWriterT)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Typeable
import Language.Fiddle.Ast
import Language.Fiddle.Ast.FileInterface (ResolvedImport)
import Language.Fiddle.Compiler
import Language.Fiddle.Compiler.Expansion
import Language.Fiddle.Internal.UnitInterface
import Language.Fiddle.Types
import Options.Applicative
import System.IO (hPutStrLn, stderr)
import Text.Printf (printf)

newtype Flags = Flags
  { importDirectories :: [FilePath]
  }

parseFlags :: Parser Flags
parseFlags =
  Flags
    <$> many
      ( strOption
          ( long "import"
              <> short 'I'
              <> metavar "DIRECTORY"
              <> help "Directory to add to the import search path"
          )
      )

importResolutionPhase ::
  CompilationPhase CurrentStage ImportsResolved
importResolutionPhase =
  CompilationPhase parseFlags getImportResolutionState resolveImports

type GlobalState = ()

type LocalState = ResolvedImports

type M = Compile GlobalState

type Annot = Commented SourceSpan

data ImportError = ImportError Text (Maybe SourceSpan)
  deriving (Show)

newtype ResolvedImports = ResolvedImports
  { importMap :: Map Text (Either ImportError 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 _ = advanceStage

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 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

deriving instance AdvanceStage CurrentStage RegisterBitsTypeRef

deriving instance AdvanceStage CurrentStage ObjType

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

deriving instance AdvanceStage CurrentStage FiddleDecl

diagnosticError :: String -> Annot -> Compile a ()
diagnosticError str a = tell [Diagnostic Error str (unCommented a)]

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
        diagnosticError "Failed to lookup imports (This is a bug)" a
        return empty
      Just (Left err) -> do
        diagnosticError (printf "Error in import %s: %s" path (show err)) a
        return empty
      Just (Right val) -> return val

    return $ ImportStatement path list v a

getImportResolutionState ::
  Flags ->
  FiddleUnit CurrentStage Identity Annot ->
  IO ResolvedImports
getImportResolutionState flags unit =
  execWriterT $
    walk doWalk unit ()
  where
    -- doWalk :: forall t'. (Walk t', Typeable t') => t' Identity Annot -> () -> WriterT ResolvedImports IO ()
    doWalk u () =
      case () of
        () | Just (ImportStatement {importPath = path}) <- castTS u -> do
          lift $ hPutStrLn stderr $ "Import path: " ++ show path
          (return () :: WriterT ResolvedImports IO ())
        _ -> return ()

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