{-# LANGUAGE ApplicativeDo #-}

module Descript.Sugar.Refine
  ( refineDDepd
  , refine
  , refineInputValIn
  , refineOutputValIn
  ) where

import qualified Descript.Sugar.Data.Value.Reg as Reg
import qualified Descript.Sugar.Data.Value.In as In
import qualified Descript.Sugar.Data.Value.Out as Out
import Descript.Sugar.Data
import qualified Descript.BasicInj.Data.Value.Reg as BasicInj.Reg
import qualified Descript.BasicInj.Data.Value.In as BasicInj.In
import qualified Descript.BasicInj.Data.Value.Out as BasicInj.Out
import qualified Descript.BasicInj.Data.Value.Gen as BasicInj.Record
import qualified Descript.BasicInj.Data.Type as BasicInj.RecordType
import qualified Descript.BasicInj.Data as BasicInj
import Descript.Misc
import Data.Monoid
import Data.Foldable
import Data.Maybe
import Core.Data.List
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NonEmpty
import Control.Monad
import Prelude hiding (mod)

-- | Expands syntactic sugar in the source using the dependency.
-- Passes dependency failures along.
refineDDepd :: (TaintAnn an)
            => BasicInj.DirtyDepd Source an
            -> BasicInj.DirtyDepd BasicInj.Source an
refineDDepd (Depd dextra x) = Depd dextra $ refine (dirtyVal dextra) x

-- | Expands syntactic sugar in the source.
refine :: (TaintAnn an)
       => BasicInj.Dep
       -> Source an
       -> BasicInj.Source an
refine extra (SourceModule mod)
  = BasicInj.SourceModule $ refineBModule extra mod
refine extra (SourceProgram prog)
  = BasicInj.SourceProgram $ refineProgram extra prog

refineProgram :: (TaintAnn an)
              => BasicInj.Dep
              -> Program an
              -> BasicInj.Program an
refineProgram extra (Program ann mod query')
  = BasicInj.Program ann mod' $ refineQueryIn scope ctx query'
  where scope = BasicInj.moduleScope mod'
        ctx = recordCtx_ <> extraCtx
        extraCtx = BasicInj.recordCtx extra
        recordCtx_ = remAnns $ BasicInj.recordCtx $ BasicInj.amodule mod'
        mod' = refineBModule extra mod

refineQueryIn :: (TaintAnn an)
              => AbsScope
              -> BasicInj.RecordCtx ()
              -> Query an
              -> BasicInj.Query an
refineQueryIn scope ctx (Query ann val)
  = BasicInj.Query ann $ refineRegValueIn scope ctx val

refineBModule :: (TaintAnn an)
              => BasicInj.Dep
              -> BModule an
              -> BasicInj.BModule an
refineBModule extra (BModule ann ictx amod)
  = BasicInj.BModule ann ictx $ refineAModuleIn scope ctx amod
  where scope = moduleDeclScope $ moduleDecl ictx
        ctx = BasicInj.recordCtx extra

refineAModuleIn :: (TaintAnn an)
                => AbsScope
                -> BasicInj.RecordCtx ()
                -> AModule an
                -> BasicInj.AModule an
refineAModuleIn scope extraCtx (AModule ann recordCtx' reduceCtx')
  = BasicInj.AModule ann recordCtx'' reduceCtx''
  where recordCtx'' = refineRecordCtxIn scope recordCtx'
        reduceCtx'' = refineReduceCtxIn scope ctx reduceCtx'
        ctx = recordCtx_ <> extraCtx
        recordCtx_ = remAnns recordCtx''

refineRecordCtxIn :: (TaintAnn an)
                  => AbsScope
                  -> RecordCtx an
                  -> BasicInj.RecordCtx an
refineRecordCtxIn scope (RecordCtx ann decls)
  = BasicInj.RecordCtx ann $ map (refineRecordDeclIn scope) decls

refineRecordDeclIn :: (TaintAnn an)
                   => AbsScope
                   -> RecordDecl an
                   -> BasicInj.RecordDecl an
refineRecordDeclIn scope (RecordDecl ann rtype)
  = BasicInj.RecordDecl ann $ refineRecordTypeIn scope rtype

refineRecordTypeIn :: (TaintAnn an)
                   => AbsScope
                   -> RecordType an
                   -> BasicInj.RecordType an
refineRecordTypeIn scope (RecordType ann head' props)
  = BasicInj.RecordType ann (FSymbol scope head') props

refineReduceCtxIn :: (TaintAnn an)
                  => AbsScope
                  -> BasicInj.RecordCtx ()
                  -> ReduceCtx an
                  -> BasicInj.ReduceCtx an
refineReduceCtxIn scope ctx (ReduceCtx ann (p :| ps))
  = BasicInj.ReduceCtx ann topPhase
  $ NonEmpty.map (refinePhaseCtxIn scope ctx) $ p :| ps
  where topPhase = BasicInj.PhaseCtx (preInsertAnn $ getAnn p) []

refinePhaseCtxIn :: (TaintAnn an)
                 => AbsScope
                 -> BasicInj.RecordCtx ()
                 -> PhaseCtx an
                 -> BasicInj.PhaseCtx an
refinePhaseCtxIn scope ctx (PhaseCtx ann reducers)
  = BasicInj.PhaseCtx ann $ map (refineReducerIn scope ctx) reducers

refineReducerIn :: (TaintAnn an)
                => AbsScope
                -> BasicInj.RecordCtx ()
                -> Reducer an
                -> BasicInj.Reducer an
refineReducerIn scope ctx (Reducer ann input' output')
  = BasicInj.Reducer ann input'' (refineOutputValIn scope ctx input_ output')
  where input_ = remAnns input''
        input'' = refineInputValIn scope ctx input'

refineRegValueIn :: (TaintAnn an)
                 => AbsScope
                 -> BasicInj.RecordCtx ()
                 -> Reg.Value an
                 -> BasicInj.Reg.Value an
refineRegValueIn scope ctx (Value ann parts)
  = BasicInj.Value ann $ map (refineRegPartIn scope ctx) parts

refineInputValIn :: (TaintAnn an)
                 => AbsScope
                 -> BasicInj.RecordCtx ()
                 -> In.Value an
                 -> BasicInj.In.Value an
refineInputValIn scope ctx (Value ann parts)
  = BasicInj.Value ann $ map (refineInPartIn scope ctx) parts

refineOutputValIn :: (TaintAnn an)
                  => AbsScope
                  -> BasicInj.RecordCtx ()
                  -> BasicInj.In.Value ()
                  -> Out.Value an
                  -> BasicInj.Out.Value an
refineOutputValIn scope ctx in' (Value ann parts)
  = BasicInj.Value ann $ map (refineOutPartIn scope ctx in') parts

refineRegPartIn :: (TaintAnn an)
                => AbsScope
                -> BasicInj.RecordCtx ()
                -> Reg.Part an
                -> BasicInj.Reg.Part an
refineRegPartIn _ _ (Reg.PartPrim prim) = BasicInj.Reg.PartPrim prim
refineRegPartIn scope ctx (Reg.PartRecord record)
  = BasicInj.Reg.PartRecord $ refineRegRecordIn scope ctx record

refineInPartIn :: (TaintAnn an)
               => AbsScope
               -> BasicInj.RecordCtx ()
               -> In.Part an
               -> BasicInj.In.Part an
refineInPartIn _ _ (In.PartPrim prim) = BasicInj.In.PartPrim prim
refineInPartIn _ _ (In.PartPrimType primType) = BasicInj.In.PartPrimType primType
refineInPartIn scope ctx (In.PartRecord record)
  = BasicInj.In.PartRecord $ refineInRecordIn scope ctx record

refineOutPartIn :: (TaintAnn an)
                => AbsScope
                -> BasicInj.RecordCtx ()
                -> BasicInj.In.Value ()
                -> Out.Part an
                -> BasicInj.Out.Part an
refineOutPartIn _ _ _ (Out.PartPrim prim) = BasicInj.Out.PartPrim prim
refineOutPartIn scope ctx in' (Out.PartRecord record)
  = BasicInj.Out.PartRecord $ refineOutRecordIn scope ctx in' record
refineOutPartIn scope _ in' (Out.PartPropPath path)
  = BasicInj.Out.PartPropPath $ refinePropPathIn scope in' path
refineOutPartIn scope ctx in' (Out.PartInjApp app)
  = BasicInj.Out.PartInjApp $ refineInjAppIn scope ctx in' app

refineRegRecordIn :: (TaintAnn an)
                  => AbsScope
                  -> BasicInj.RecordCtx ()
                  -> Reg.Record an
                  -> BasicInj.Reg.Record an
refineRegRecordIn scope ctx (Record ann head' props)
  = BasicInj.Record ann head''
  $ zipWith (refineRegPropIn scope ctx head_) [0..] props
  where head_ = remAnns head''
        head'' = FSymbol scope head'

refineInRecordIn :: (TaintAnn an)
                 => AbsScope
                 -> BasicInj.RecordCtx ()
                 -> In.Record an
                 -> BasicInj.In.Record an
refineInRecordIn scope ctx (Record ann head' props)
  = BasicInj.Record ann head''
  $ zipWith (refineInPropIn scope ctx head_) [0..] props
  where head_ = remAnns head''
        head'' = FSymbol scope head'

refineOutRecordIn :: (TaintAnn an)
                  => AbsScope
                  -> BasicInj.RecordCtx ()
                  -> BasicInj.In.Value ()
                  -> Out.Record an
                  -> BasicInj.Out.Record an
refineOutRecordIn scope ctx in' (Record ann head' props)
  = BasicInj.Record ann head''
  $ zipWith (refineOutPropIn scope ctx in' head_) [0..] props
  where head_ = remAnns head''
        head'' = FSymbol scope head'

refineRegPropIn :: (TaintAnn an)
                => AbsScope
                -> BasicInj.RecordCtx ()
                -> BasicInj.FSymbol ()
                -> Int
                -> Reg.Property an
                -> BasicInj.Reg.Property an
refineRegPropIn scope ctx head' idx (Property ann key val)
  = BasicInj.Property ann key' $ refineRegValueIn scope ctx val
  where key' = resolvePropKey ctx head' idx keyAnn `fromMaybe` key
        keyAnn = preInsertAnn ann

refineInPropIn :: (TaintAnn an)
               => AbsScope
               -> BasicInj.RecordCtx ()
               -> BasicInj.FSymbol ()
               -> Int
               -> In.Property an
               -> BasicInj.In.Property an
refineInPropIn scope ctx head' idx (Property ann key val)
  = BasicInj.Property ann key' $ refineInOptValueIn scope ctx val
  where key' = resolvePropKey ctx head' idx keyAnn `fromMaybe` key
        keyAnn = preInsertAnn ann

refineOutPropIn :: (TaintAnn an)
                => AbsScope
                -> BasicInj.RecordCtx ()
                -> BasicInj.In.Value ()
                -> BasicInj.FSymbol ()
                -> Int
                -> Out.Property an
                -> BasicInj.Out.Property an
refineOutPropIn scope ctx in' head' idx (Property ann key val)
  = BasicInj.Property ann key' $ refineOutputValIn scope ctx in' val
  where key' = resolvePropKey ctx head' idx keyAnn `fromMaybe` key
        keyAnn = preInsertAnn ann

refineInOptValueIn :: (TaintAnn an)
                   => AbsScope
                   -> BasicInj.RecordCtx ()
                   -> In.OptValue an
                   -> BasicInj.In.OptValue an
refineInOptValueIn _ _ In.NothingValue = BasicInj.In.NothingValue
refineInOptValueIn scope ctx (In.JustValue val)
  = BasicInj.In.JustValue $ refineInputValIn scope ctx val

refinePropPathIn :: (TaintAnn an)
                 => AbsScope
                 -> BasicInj.In.Value ()
                 -> PropPath an
                 -> BasicInj.PropPath an
refinePropPathIn scope in' (PropPath ann (x :| xs))
  = BasicInj.PropPath ann $ x' :| refineSubPathIn scope inSub xs
  where inSub = inputInElem x_ in'
        x_ = remAnns x'
        x' = refinePathElemIn scope in' x

refineSubPathIn :: (TaintAnn an)
                => AbsScope
                -> BasicInj.In.Value ()
                -> SubPropPath an
                -> BasicInj.SubPropPath an
refineSubPathIn _ _ [] = []
refineSubPathIn scope in' (x : xs)
  = x' : refineSubPathIn scope inSub xs
  where inSub = inputInElem x_ in'
        x_ = remAnns x'
        x' = refinePathElemIn scope in' x

refinePathElemIn :: (TaintAnn an)
                 => AbsScope
                 -> BasicInj.In.Value ()
                 -> PathElem an
                 -> BasicInj.PathElem an
refinePathElemIn _ in' (PathElemImp key) = BasicInj.PathElem ann key head'
  where head' = resolvePathHead in' key
        ann = getAnn key
refinePathElemIn scope _ (PathElemExp ann key head')
  = BasicInj.PathElem ann key $ FSymbol scope head'

refineInjAppIn :: (TaintAnn an)
               => AbsScope
               -> BasicInj.RecordCtx ()
               -> BasicInj.In.Value ()
               -> Out.InjApp an
               -> BasicInj.Out.InjApp an
refineInjAppIn scope ctx in' (Out.InjApp ann funcId' params')
  = BasicInj.Out.InjApp ann funcId'
  $ map (refineOutParamIn scope ctx in') params'

refineOutParamIn :: (TaintAnn an)
                 => AbsScope
                 -> BasicInj.RecordCtx ()
                 -> BasicInj.In.Value ()
                 -> Out.InjParam an
                 -> BasicInj.Out.InjParam an
refineOutParamIn scope ctx in' (Out.InjParam ann val)
  = BasicInj.Out.InjParam ann $ refineOutputValIn scope ctx in' val

inputInElem :: BasicInj.PathElem ()
            -> BasicInj.In.Value ()
            -> BasicInj.In.Value ()
inputInElem (BasicInj.PathElem () key head')
  = fold
  . ( BasicInj.In.optValToMaybeVal
  <=< BasicInj.lookupProp key
  <=< BasicInj.recWithHead head' )

resolvePropKey :: (TaintAnn an)
               => BasicInj.RecordCtx ()
               -> BasicInj.FSymbol ()
               -> Int
               -> an
               -> BasicInj.Symbol an
resolvePropKey ctx head' idx ann
  = ann <$ (undefinedSym `fromMaybe` resolvePropKey_ ctx head' idx)

resolvePropKey_ :: BasicInj.RecordCtx ()
                -> BasicInj.FSymbol ()
                -> Int
                -> Maybe (BasicInj.Symbol ())
resolvePropKey_ ctx head' idx
    = (!? idx) . BasicInj.RecordType.properties
  =<< BasicInj.lookupRecordType head' ctx

resolvePathHead :: (TaintAnn an)
                => BasicInj.In.Value ()
                -> Symbol an
                -> FSymbol an
resolvePathHead in' key
  = ann' <$ (BasicInj.undefinedFSym `fromMaybe` resolvePathHead_ in' key_)
  where key_ = remAnns key
        ann' = postInsertAnn $ getAnn key

resolvePathHead_ :: BasicInj.In.Value ()
                 -> Symbol ()
                 -> Maybe (FSymbol ())
resolvePathHead_ (BasicInj.Value () parts) key
  = case filter (BasicInj.recHasProp key) $ mapMaybe BasicInj.partToRec parts of
         [record] -> Just $ BasicInj.Record.head record
         _ -> Nothing