{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TupleSections #-}

module Descript.BasicInj.Process.Reduce.PropTrans
  ( PropTranses
  , PropTrans (..)
  , apPropTranses
  , subPropTranses
  ) where

import qualified Descript.BasicInj.Data.Value.In as In
import qualified Descript.BasicInj.Data.Value.Out as Out
import Descript.BasicInj.Data
import Descript.Misc
import Data.Maybe
import Data.List
import Data.List.NonEmpty (NonEmpty (..))
import Control.Applicative

-- | Replaces properties within a value.
type PropTranses = [PropTrans]

-- | Replaces every occurrence of a property path with a union of the
-- new subpaths. If a subpath is empty, it's replaced by the entire
-- input value with immediate properties.
data PropTrans
  = PropTrans
  { propTransOld :: PropPath ()
  , propTransNews :: [SubPropPath ()]
  } deriving (Show)

-- | Merges the transformations, then applies them all.
apPropTranses :: (TaintAnn an)
              => In.Value ()
              -> PropTranses
              -> Out.Value an
              -> Out.Value an
apPropTranses in' transes x
  = foldl' (flip $ apPropTrans in') x $ mergeTranses transes

-- | Adds outputs from more general transes to more specific ones.
-- For example:
--
-- >>> mergeTranses ['a>c>d -> ', 'a -> z', 'a>b -> y']
-- ['a>c>d -> z>c>d', 'a -> z', 'a>b -> y | z>b']
mergeTranses :: PropTranses -> PropTranses
mergeTranses = foldl' (flip mergeAddTrans) []

-- | Adds outputs from more general transes to the new trans, and adds
-- outputs from the new trans to more specific ones.
mergeAddTrans :: PropTrans -> PropTranses -> PropTranses
mergeAddTrans x [] = [x]
mergeAddTrans x (y : ys) = y' : mergeAddTrans x' ys
  where (x', y') = tryMergeTrans x y

-- | If one of the transes is a prefix of the other, returns both so
-- the longer one has the smaller one's outputs. Otherwise returns both
-- unaffected.
tryMergeTrans :: PropTrans -> PropTrans -> (PropTrans, PropTrans)
tryMergeTrans x y = (x, y) `fromMaybe` mergeTrans x y

-- | If one of the transes is a prefix of the other, returns both so
-- the longer one has the smaller one's outputs.
mergeTrans :: PropTrans -> PropTrans -> Maybe (PropTrans, PropTrans)
mergeTrans x y
    = (x, ) <$> mergeTrans1Way x y
  <|> (, y) <$> mergeTrans1Way y x

-- | If the first transes is a prefix of the second, returns the second
-- so it includes the first. Otherwise returns 'Nothing'.
mergeTrans1Way :: PropTrans -> PropTrans -> Maybe PropTrans
mergeTrans1Way (PropTrans xOld xNews) (PropTrans yOld yNews)
  = case xOld `stripPrefixPath` yOld of
         Nothing -> Nothing
         Just yRest -> Just $ PropTrans yOld $ yNews ++ xNews'
           where xNews' = map (++ yRest) xNews

{-
  = case old `stripPrefixPath` path_ of
         Nothing -> singletonValue $ Out.PartPropPath path
         Just suf_ -> ann' <$ mconcat (map (subPathVal in' . (++ suf_)) news)
  where path_ = remAnns path
        ann' = taint $ getAnn path
-}

-- | Assumes there is a transformation for every reasonable path, and
-- the transformations were all merged (more general transformation
-- outputs, specified, were added to more specific outputs).
apPropTrans :: (TaintAnn an)
            => In.Value ()
            -> PropTrans
            -> Out.Value an
            -> Out.Value an
apPropTrans in' trans (Value ann parts)
  = reconValue ann parts $ map (apPropTransToPart in' trans) parts

apPropTransToPart :: (TaintAnn an)
                  => In.Value ()
                  -> PropTrans
                  -> Out.Part an
                  -> Out.Value an
apPropTransToPart _ _ (Out.PartPrim prim) = singletonValue $ Out.PartPrim prim
apPropTransToPart in' trans (Out.PartRecord record)
  = singletonValue $ Out.PartRecord $ apPropTransToRecord in' trans record
apPropTransToPart in' trans (Out.PartPropPath path)
  = apPropTransToPath in' trans path
apPropTransToPart in' trans (Out.PartInjApp app)
  = singletonValue $ Out.PartInjApp $ apPropTransToInjApp in' trans app

apPropTransToRecord :: (TaintAnn an)
                    => In.Value ()
                    -> PropTrans
                    -> Out.Record an
                    -> Out.Record an
apPropTransToRecord in' trans (Record ann head' props)
  -- No need for reRecord because prop counts, and thus annotation, will
  -- always stay the same.
  = Record ann head' $ map (apPropTransToProp in' trans) props

apPropTransToProp :: (TaintAnn an)
                  => In.Value ()
                  -> PropTrans
                  -> Out.Property an
                  -> Out.Property an
apPropTransToProp in' trans (Property ann key val)
  -- No need for reProperty because value printability, and thus
  -- annotation, will always stay the same.
  = Property ann key $ apPropTrans in' trans val

apPropTransToInjApp :: (TaintAnn an)
                    => In.Value ()
                    -> PropTrans
                    -> Out.InjApp an
                    -> Out.InjApp an
apPropTransToInjApp in' trans (Out.InjApp ann funcId' params')
  = Out.InjApp ann funcId' $ map (apPropTransToInjParam in' trans) params'

apPropTransToInjParam :: (TaintAnn an)
                      => In.Value ()
                      -> PropTrans
                      -> Out.InjParam an
                      -> Out.InjParam an
apPropTransToInjParam in' trans (Out.InjParam ann val)
  = Out.InjParam ann $ apPropTrans in' trans val

apPropTransToPath :: (TaintAnn an)
                  => In.Value ()
                  -> PropTrans
                  -> PropPath an
                  -> Out.Value an
apPropTransToPath in' (PropTrans old news) path
-- Assumes there is a transformation for every reasonable path, and the
-- transformations were all merged. Otherwise would need to check prefix
-- instead of full path.
  | path /@= old = singletonValue $ Out.PartPropPath path
  | otherwise = ann' <$ mconcat (map (subPathVal in') news)
  where ann' = taint $ getAnn path

-- | An output value which refers to the given sub-path, given the
-- corresponding input. If the sub-path is empty, the value will refer
-- to the entire input (via 'fullOut'). Otherwise it will just contain a
-- single property path.
subPathVal :: In.Value () -> SubPropPath () -> Out.Value ()
subPathVal in' [] = fullOut in'
subPathVal _ (x : xs) = singletonValue $ Out.PartPropPath $ PropPath () $ x :| xs

-- | This output will re-produce everything which was consumed by the
-- input. It's semantically equivalent to an empty property path, but
-- those don't exist. Note that a top-level primitive type has no full
-- output (you can't re-produce it), so that will raise an error.
fullOut :: In.Value () -> Out.Value ()
fullOut (Value () parts) = Value () $ map fullOutPart parts

fullOutPart :: In.Part () -> Out.Part ()
fullOutPart (In.PartPrim prim) = Out.PartPrim prim
fullOutPart (In.PartPrimType _)
  = error "Top-level primitive type has no full output - you can't \
          \reproduce it in an output value, because it's not a single \
          \value and it doesn't correspond to a property path."
fullOutPart (In.PartRecord record) = Out.PartRecord $ fullOutRecord record

fullOutRecord :: In.Record () -> Out.Record ()
fullOutRecord (Record () head' props)
  = Record () head' $ map (fullOutProp head') props

fullOutProp :: FSymbol () -> In.Property () -> Out.Property ()
fullOutProp head' (Property () key _) = Property () key $ Out.immPathVal elem'
  where elem' = PathElem () key head'

-- | Prepends the element to all property paths in all transformations.
subPropTranses :: PathElem () -> PropTranses -> PropTranses
subPropTranses = map . subPropTrans

-- | Prepends the element to the input and all output paths.
subPropTrans :: PathElem () -> PropTrans -> PropTrans
subPropTrans x (PropTrans old news)
  = PropTrans
  { propTransOld = subPath x old
  , propTransNews = map (x :) news
  }