{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns #-}

-- Things used by both RawToHaskell and PureToHaskell.
module Trans.ToHaskellCommon where

import qualified Capnp as Capnp
import qualified Capnp.Basics as B
import Capnp.Classes (encode)
import Capnp.Repr.Parsed (Parsed)
import qualified Data.ByteString.Lazy as LBS
import Data.Char (toUpper)
import Data.Maybe (fromJust)
import qualified Data.Set as S
import qualified Data.Text as T
import Data.Word
import GHC.Exts (fromList)
import qualified IR.Common as C
import IR.Haskell
import qualified IR.Name as Name
import System.FilePath (splitDirectories)
import Text.Printf (printf)

-- Misc shortcuts for common constructs:

std_ :: Name.UnQ -> Name.GlobalQ
std_ name = gName ["Std_"] (Name.mkLocal Name.emptyNS name)

eStd_ :: Name.UnQ -> Exp
eStd_ = EGName . std_

ePureUnit :: Exp
ePureUnit = EApp (eStd_ "pure") [ETup []]

tStd_ :: Name.UnQ -> Type
tStd_ = TGName . std_

gName :: [T.Text] -> Name.LocalQ -> Name.GlobalQ
gName parts local =
  Name.GlobalQ
    { globalNS = Name.NS parts,
      local
    }

egName :: [T.Text] -> Name.LocalQ -> Exp
egName parts local = EGName $ gName parts local

euName :: Name.UnQ -> Exp
euName = ELName . Name.mkLocal Name.emptyNS

tgName :: [T.Text] -> Name.LocalQ -> Type
tgName parts local = TGName $ gName parts local

pgName :: [T.Text] -> Name.LocalQ -> [Pattern] -> Pattern
pgName parts local = PGCtor (gName parts local)

puName :: Name.UnQ -> [Pattern] -> Pattern
puName = PLCtor . Name.unQToLocal

tuName :: Name.UnQ -> Type
tuName = TLName . Name.unQToLocal

iValue :: Name.UnQ -> [Pattern] -> Exp -> InstanceDef
iValue name params value = IdValue DfValue {name, params, value}

iType :: Name.UnQ -> [Type] -> Type -> InstanceDef
iType name params value = IdType $ TypeAlias name params value

readCtx :: T.Text -> T.Text -> Type
readCtx m msg =
  TApp
    (tgName ["Untyped"] "ReadCtx")
    [ TVar m,
      TVar msg
    ]

rwCtx :: T.Text -> T.Text -> Type
rwCtx m s =
  TApp
    (tgName ["Untyped"] "RWCtx")
    [ TVar m,
      TVar s
    ]

eGetWordField :: Exp -> C.DataLoc -> Exp
eGetWordField struct C.DataLoc {dataIdx, dataOff, dataDef} =
  EApp
    (egName ["GenHelpers"] "getWordField")
    [ struct,
      EInt $ fromIntegral dataIdx,
      EInt $ fromIntegral dataOff,
      EInt $ fromIntegral dataDef
    ]

idToModule :: Word64 -> [Name.UnQ]
idToModule fileId =
  ["Capnp", "Gen", "ById", Name.UnQ $ T.pack $ printf "X%x" fileId]

instance_ :: [Type] -> [T.Text] -> Name.LocalQ -> [Type] -> [InstanceDef] -> Decl
instance_ ctx [] className tys defs =
  DcInstance
    { ctx,
      typ = TApp (TLName className) tys,
      defs
    }
instance_ ctx classNS className tys defs =
  DcInstance
    { ctx,
      typ = TApp (tgName classNS className) tys,
      defs
    }

-- | Transform the file path into a valid haskell module name.
-- TODO: this is a best-effort transformation; it gives good
-- results on the schema I(zenhack) have found in the wild, but
-- may fail to generate valid/non-overlapping module names in
-- all cases.
--
-- This generates the bit that is unique to the specific file
-- name and common to both raw and pure backends, so e.g. for
-- @myorg/example.capnp@ it generates @["Myorg", "Example"]@.
makeModName :: FilePath -> [Name.UnQ]
makeModName fileName =
  [Name.UnQ (T.pack (mangleSegment seg)) | seg <- splitDirectories fileName]
  where
    mangleSegment "c++.capnp" = "Cxx"
    mangleSegment "" = error "Unexpected empty file name"
    mangleSegment (c : cs) = go (toUpper c : cs)
      where
        go ('-' : c : cs) = toUpper c : go cs
        go ".capnp" = ""
        go [] = ""
        go (c : cs) = c : go cs

-- | Fix the capnp imports of a module, so that they correspond to the
-- imports actually used in the body of the module and/or export list.
--
-- Note that this only looks at imports of the form Capnp.Gen....; other
-- imports are not touched.
fixImports :: Module -> Module
fixImports m@Module {modImports} =
  let namespaces =
        S.toList $
          S.fromList -- deduplicate
            [ nsParts
              | Name.GlobalQ
                  { globalNS = Name.NS nsParts@(map T.unpack -> "Capnp" : "Gen" : _)
                  } <-
                  S.toList (findGNames m)
            ]
      neededImports =
        [ ImportQual {parts = map Name.UnQ nsParts}
          | nsParts <- namespaces
        ]
   in m {modImports = modImports ++ neededImports}

toTVars :: [Name.UnQ] -> [Type]
toTVars = map (TVar . Name.typeVarName)

makePtrBytes :: Maybe (Parsed B.AnyPointer) -> LBS.ByteString
makePtrBytes ptr =
  Capnp.msgToLBS $ fromJust $ Capnp.createPure Capnp.defaultLimit $ do
    msg <- Capnp.newMessage Nothing
    rootPtr <-
      encode msg $
        B.Struct
          (fromList [])
          (fromList [ptr])
    Capnp.setRoot rootPtr
    pure msg

class HasGNames a where
  -- | Collect all of the 'Name.GlobalQ's used in the module.
  --
  -- This seems like it would be the perfect use case for something
  -- like syb or similar libraries, but I(zenhack) haven't taken the
  -- time to fully wrap my head around how to use those yet, so we do
  -- it the boilerplate-heavy way.
  findGNames :: a -> S.Set Name.GlobalQ

instance HasGNames Module where
  findGNames Module {modExports = Just exports, modDecls} =
    S.unions $ map findGNames exports ++ map findGNames modDecls
  findGNames Module {modExports = Nothing, modDecls} =
    S.unions $ map findGNames modDecls

instance HasGNames Export where
  findGNames (ExportGCtors name) = S.singleton name
  findGNames (ExportGName name) = S.singleton name
  findGNames _ = S.empty

instance HasGNames Decl where
  findGNames (DcData d) = findGNames d
  findGNames DcValue {typ, def} = findGNames typ `S.union` findGNames def
  findGNames DcInstance {ctx, typ, defs} =
    S.unions
      [ S.unions $ map findGNames ctx,
        findGNames typ,
        S.unions $ map findGNames defs
      ]
  findGNames (DcTypeInstance alias orig) = findGNames alias `S.union` findGNames orig
  findGNames (DcDeriveInstance ctx typ) = findGNames (TCtx ctx typ)
  findGNames DcClass {ctx, decls} =
    S.unions $ map findGNames ctx ++ map findGNames decls

instance HasGNames DataDecl where
  findGNames Data {typeArgs, dataVariants} =
    S.unions $ map findGNames typeArgs ++ map findGNames dataVariants

instance HasGNames ClassDecl where
  findGNames (CdValueDecl _ ty) = findGNames ty
  findGNames (CdValueDef d) = findGNames d
  findGNames (CdMinimal _) = S.empty

instance HasGNames InstanceDef where
  findGNames (IdValue d) = findGNames d
  findGNames (IdData d) = findGNames d
  findGNames (IdType t) = findGNames t

instance HasGNames TypeAlias where
  findGNames (TypeAlias _ ts t) = S.unions $ map findGNames (t : ts)

instance HasGNames ValueDef where
  findGNames DfValue {value, params} =
    S.unions $ findGNames value : map findGNames params

instance HasGNames DataVariant where
  findGNames DataVariant {dvArgs} = findGNames dvArgs

instance HasGNames DataArgs where
  findGNames (APos tys) = S.unions $ map findGNames tys
  findGNames (ARec fields) = S.unions $ map (findGNames . snd) fields

instance HasGNames Type where
  findGNames (TGName n) = S.singleton n
  findGNames (TApp t ts) = S.unions $ map findGNames (t : ts)
  findGNames (TFn ts) = S.unions $ map findGNames ts
  findGNames (TCtx ts t) = S.unions $ map findGNames (t : ts)
  findGNames _ = S.empty

instance HasGNames Exp where
  findGNames (EApp e es) = S.unions $ map findGNames (e : es)
  findGNames (EFApp e es) = S.unions $ map findGNames (e : es)
  findGNames (EDo ds e) = S.unions $ findGNames e : map findGNames ds
  findGNames (EBind x y) = findGNames x `S.union` findGNames y
  findGNames (ETup es) = S.unions $ map findGNames es
  findGNames (ECase e arms) =
    S.unions
      [ findGNames e,
        S.unions $ map (findGNames . fst) arms,
        S.unions $ map (findGNames . snd) arms
      ]
  findGNames (ETypeAnno e t) = findGNames e `S.union` findGNames t
  findGNames (ELambda ps e) = S.unions $ findGNames e : map findGNames ps
  findGNames (ERecord e fields) = S.unions $ findGNames e : map (findGNames . snd) fields
  findGNames _ = S.empty

instance HasGNames Do where
  findGNames (DoBind _ e) = findGNames e
  findGNames (DoE e) = findGNames e

instance HasGNames Pattern where
  findGNames (PLCtor _ ps) = S.unions $ map findGNames ps
  findGNames (PGCtor n ps) = S.unions $ S.singleton n : map findGNames ps
  findGNames _ = S.empty