{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
module Rename.Rename where

import Ahc.Data.Lens

import Control.Monad.IO.Class
import Control.Monad.Reader
import Control.Concurrent
import Control.Exception

import qualified Data.HashMap.Strict as HashMap
import Data.HashMap.Strict (HashMap)
import Data.ByteString (ByteString)
import Data.Text (Text)

import Debug.Trace

import Development.Shake ( Action, askOracle )
import Development.Shake.Classes

import Errors

import qualified Frontend.Syntax as Fe
import Frontend.Parser.Posn
import Frontend.Syntax.Var
import Frontend.Syntax

import GHC.Generics

import Main.Queries

import Rename.Types

import Text.Show.Pretty
import Control.Applicative

newtype Rename a = RenameM { runRename :: RenameCtx -> Action a }
  deriving (Functor, Applicative, Monad, MonadIO, MonadReader RenameCtx)
    via ReaderT RenameCtx Action

liftAct :: Action a -> Rename a
liftAct a = RenameM $ \x -> a

data RenameCtx =
  RenameCtx
  { rcOurModId    :: ModuleId
  , rcOurFilePath :: String
  , rcScope       :: Scope
  , rcThisModule  :: MVar Scope
  , rcSpanStart   :: Posn
  , rcSpanEnd     :: Posn
  }

_rcScope :: Lens' RenameCtx Scope
_rcScope = lens rcScope (\s x -> s { rcScope = x })

renamerSourcePos :: Posn -> Posn -> RenameCtx -> RenameCtx
renamerSourcePos start end x = x { rcSpanStart = start, rcSpanEnd = end }

renameModule :: FeModule ParsedVar -> Rename RenamedMod
renameModule Fe.Module{moduleExports, moduleImports, moduleItems, Fe.moduleName} = do
  (imports, scopes) <- fmap unzip $ traverse importModule moduleImports
  our_mod_id <- asks rcOurModId
  local (extendBy (mconcat scopes)) $ do
    items <- renameModuleBody moduleItems

    this_mod <- liftIO . readMVar =<< asks rcThisModule
    liftIO $ print this_mod
    export <- traverse (renameModVarList this_mod) moduleExports

    our_mod_name <- makeVariable Rename.Types.Module moduleName
    let
      scope = restrictToModVarList export this_mod
      rnm_mod = Fe.Module { Fe.moduleName = our_mod_name
                          , moduleExports = export
                          , moduleImports = imports
                          , moduleItems   = items
                          }
    pure (RnmMod rnm_mod scope)

restrictToModVarList :: Maybe [NamespacedItem FqVar] -> Scope -> Scope
restrictToModVarList Nothing x = x
restrictToModVarList (Just xs) sc = foldr del sc xs where
  del :: NamespacedItem FqVar -> Scope -> Scope
  del (IEVar n)    = over scopeNames (HashMap.delete (fqVarName n))
  del (IECon n)    = over scopeNames (HashMap.delete (fqVarName n)) . over scopeTypes (HashMap.delete (fqVarName n))
  del (IEModule n) = over scopeNamespaces (HashMap.delete (fqVarName n))

renameModVarList :: Scope -> [NamespacedItem ParsedVar] -> Rename [NamespacedItem FqVar]
renameModVarList sc (x:xs) =
  case x of
    IEVar pv    ->
      case HashMap.lookup (varId pv) (sc ^. scopeNames) of
        Just [t] -> (IEVar t:) <$> renameModVarList sc xs
        _ -> notInScope pv
    IECon pv    ->
      case HashMap.lookup (varId pv) (sc ^. scopeNames) <|> HashMap.lookup (varId pv) (sc ^. scopeTypes) of
        Just [t] -> (IECon t:) <$> renameModVarList sc xs
        _ -> notInScope pv
    IEModule pv ->
      case HashMap.lookup (varId pv) (sc ^. scopeNamespaces) of
        Just (t, _) -> (IEModule t:) <$> renameModVarList sc xs
        _ -> notInScope pv
renameModVarList sc [] = pure []

renameModuleBody :: [ModuleItem ParsedVar] -> Rename [ModuleItem FqVar]
renameModuleBody = wrap where
  wrap [] = pure []
  wrap xs = do
    var <- asks rcThisModule
    go xs

  go [] = pure []
  go (ModDecl (TySig vars tipe begin end) begin' end':items) = local (renamerSourcePos begin end) $ do
    tipe <- renameTypeToplevel tipe
    vars <- traverse (makeVariable_unique Variable scopeNames) vars
    r <- go items
    insertVariables scopeNames vars $
      pure (ModDecl (TySig (map snd vars) tipe begin end) begin' end':r)
  go (x:items) = error $ "show item: " ++ ppShow x

renameTypeToplevel :: FeType ParsedVar -> Rename (FeType FqVar)
renameTypeToplevel (Tytup xs) = Tytup <$> traverse renameTypeToplevel xs
renameTypeToplevel (SPType ty s e) =
  SPType <$> local (renamerSourcePos s e) (renameTypeToplevel ty)
         <*> pure s
         <*> pure e
-- rename units to units
renameTypeToplevel (Tycon (BuiltinId _ (BuiltinTuple 0) _ _)) = pure (Tytup [])
renameTypeToplevel (Tycon x) = do
  var <- useVariable scopeTypes x
  pure (Tycon var)
renameTypeToplevel x = error $ "idk how to rename: " ++ show x

extendBy :: Scope -> RenameCtx -> RenameCtx
extendBy s rc = rc{rcScope = s <> rcScope rc}

importModule :: ModuleImport ParsedVar -> Rename (ModuleImport FqVar, Scope)
importModule Import{importMod, importList, importQualified, importAlias, importBegin, importEnd} = do
  fp <- asks rcOurFilePath
  our_mod_id <- asks rcOurModId
  res_fp <- liftAct $ askOracle $
    AhcModuleFilepath { ahcfpImportingModule = fp
                      , ahcfpModName = importMod
                      }

  renamed_import <- liftAct $ askOracle $ AhcRenamedModule res_fp
  let
    orig = rnmModSignature renamed_import
    other_mod = Fe.moduleName $ rnmModModule renamed_import
  il <- traverse (renameModVarList orig) importList
  orig <- pure $ restrictToModVarList il orig

  mod <- pure $
    if importQualified
      then mempty { _scopeNamespaces = HashMap.singleton (varId importMod) (other_mod, orig) }
      else orig

  mod <- pure $
    case importAlias of
      Just x -> mod & scopeNamespaces %~ HashMap.insert (varId x) (other_mod, orig)
      Nothing -> mod

  our_mod_name <- makeVariable Rename.Types.Module importMod
  alias <- traverse (makeVariable Rename.Types.Module) importAlias
  let
    imp = Import our_mod_name il importQualified alias importBegin importEnd

  pure (imp, mod)

makeVariable :: MonadReader RenameCtx m => Namespace -> ParsedVar -> m FqVar
makeVariable nsc x = do
  context <- ask
  pure $
    FqVar { fqVarName      = varId x
          , fqVarModule    = rcOurModId context
          , fqVarNamespace = nsc
          , fqBegin        = rcSpanStart context
          , fqEnd          = rcSpanEnd context
          }

makeVariable_unique :: Namespace -> Lens' Scope (HashMap Text [FqVar]) -> ParsedVar -> Rename (Text, FqVar)
makeVariable_unique this_ns ns x = do
  var <- asks rcThisModule
  scope <- liftIO . readMVar $ var
  let
    the_hm = scope ^. ns
  case HashMap.lookup (varId x) the_hm of
    Just (orig:_)  -> redeclaration orig x
    _ -> do
      new_var <- makeVariable this_ns x
      liftIO $ modifyMVar_ var $ \s -> pure $ over ns (insertVariable (varId x) new_var) s
      pure (varId x, new_var)

insertVariable :: Text -> FqVar -> HashMap Text [FqVar] -> HashMap Text [FqVar]
insertVariable c v = HashMap.alter go c where
  go (Just xs) = Just (v:xs)
  go Nothing = Just [v]

insertVariables :: Lens' Scope (HashMap Text [FqVar]) -> [(Text, FqVar)] -> Rename a -> Rename a
insertVariables lens vars = local (over (_rcScope . lens) go) where
  go x = foldr (\(x, y) -> insertVariable x y) x vars

useVariable :: Lens' Scope (HashMap Text [FqVar]) -> ParsedVar -> Rename FqVar
useVariable final_ns var =
  do scope <- asks rcScope
     go scope var
  where
    go scope (QualVar id prfx begin end) =
      case HashMap.lookup prfx (scope ^. scopeNamespaces) of
        Nothing -> notInScope var
        Just (_, sc) -> local (\s -> s { rcScope = sc }) $ useVariable final_ns (UnqualVar id begin end)

    go scope (UnqualVar id begin end) =
      case HashMap.lookup id (scope ^. final_ns) of
        Just [var] -> pure var
        Just vars -> ambiguousVariables vars
        Nothing -> notInScope var

    go _ x = error $ "idk how to rename yet: " ++ show x

ambiguousVariables :: [FqVar] -> Rename a
ambiguousVariables vs = throwRenamer $ \e -> e { errorMessage = show vs }

notInScope :: HasPosn x => x -> Rename a
notInScope x =
  local (renamerSourcePos (startPosn x) (endPosn x)) $
    throwRenamer $ \e -> e { errorMessage = "variable not in scope" }

redeclaration :: FqVar -> ParsedVar -> Rename a
redeclaration first v = do
  liftIO $ pPrint first
  throwRenamer $ \e -> e { errorMessage = "variables can only have one declaration."
                         , errorInlineDesc = Just "first declaration here"
                         , errorPointers = pointTo v v "redeclaration here"
                         , errorBegin = startPosn first
                         , errorEnd = endPosn first
                         }

throwRenamer :: (MonadIO m, MonadReader RenameCtx m) => (AhcError -> AhcError) -> m b
throwRenamer cont = do
  context <- ask
  let
    fname = moduleFp (rcOurModId context)
    errorBegin = rcSpanStart context
    errorEnd = rcSpanEnd context
  liftIO . throwIO $ cont (emptyError { errorFilename = fname, errorBegin, errorEnd })