{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}

module Descript.BasicInj.Data.Value.Gen
  ( GenProperty (..)
  , GenRecord (..)
  , GenValue (..)
  , PartProperty
  , PartRecord
  , GenPropVal (..)
  , GenPart (..)
  , singletonValue
  , isEmpty
  , primParts
  , mapParts
  , foldMapParts
  , traverseParts
  , recWithHead
  , forceRecWithHead
  , deleteRecWithHead
  , mergeAddRecord
  , partHasHead
  , recHasHead
  , recHasProp
  , lookupProp
  , forceLookupProp
  , mapPropVals
  , foldMapPropVals
  , traversePropVals
  , reconValue
  , reValue
  , reRecord
  , reProperty
  ) where

import Prelude hiding (head)
import Descript.BasicInj.Data.Atom
import Descript.Misc
import Data.Monoid as M
import Core.Data.Monoid
import Data.Semigroup as S
import Core.Data.Semigroup
import Data.Proxy
import Data.Maybe
import Data.List hiding (head)
import Core.Data.List.Assoc
import qualified Data.List.NonEmpty as NonEmpty

-- | A property whose values are of type @v@.
data GenProperty v an
  = Property
  { propertyAnn :: an
  , propertyKey :: Symbol an
  , propertyValue :: v an
  } deriving (Eq, Ord, Read, Show, Functor, Foldable, Traversable)

-- | A record whose property values are of type @v@ - a record in a
-- value of type @v@.
--
-- Records are the main data of Descript. A record with no properties
-- encodes a singleton value or leaf. A record with properties encodes a
-- structure, or a tree where the properties are branches. Records are
-- similar to products, but their fields are named and they can be
-- differentiated by thier heads.
data GenRecord v an
  = Record
  { recordAnn :: an
  , head :: FSymbol an
  , properties :: [GenProperty v an]
  } deriving (Eq, Ord, Read, Show, Functor, Foldable, Traversable)

-- | A value whose parts are of type 'p'.
data GenValue p an
  = Value
  { valueAnn :: an
  , valueParts :: [p an]
  } deriving (Eq, Ord, Read, Show, Functor, Foldable, Traversable)

-- | A property for a particular part.
type PartProperty p an = GenProperty (PartPropVal p) an

-- | A record for a particular part.
type PartRecord p an = GenRecord (PartPropVal p) an

class (Semigroup1 v, Monoid1 v) => GenPropVal v where
  -- | Whether this prints in a property definition.
  -- If true, it would print @key: ...@. If false, it would print @key@.
  doesPrint :: v an -> Bool

-- | A part of a value.
class (Ann p, GenPropVal (PartPropVal p), Eq (p ())) => GenPart p where
  type PartPropVal p :: * -> *

  partToPrim :: p an -> Maybe (Prim an)
  partToRec :: p an -> Maybe (PartRecord p an)
  primToPart :: Proxy p -> Prim an -> p an
  recToPart :: Proxy p -> PartRecord p an -> p an
  mergeAddPart :: (Semigroup an) => p an -> [p an] -> [p an]

instance (GenPart p, Semigroup an, Monoid an) => Monoid (GenValue p an) where
  mempty = Value mempty []
  mappend = (S.<>)
  mconcat vals
    = Value
    { valueAnn = mconcat $ map valueAnn vals
    , valueParts = cmergeParts $ map valueParts vals
    }

instance (GenPart p, Semigroup an) => Semigroup (GenValue p an) where
  Value xAnn xParts <> Value yAnn yParts
    = Value (xAnn S.<> yAnn) $ mergeParts xParts yParts
  sconcat vals
    = Value
    { valueAnn = sconcat $ NonEmpty.map valueAnn vals
    , valueParts = cmergeParts $ map valueParts $ NonEmpty.toList vals
    }

instance (GenPart p) => GenPropVal (GenValue p) where
  doesPrint _ = True

instance (GenPart p) => Monoid1 (GenValue p)
instance (GenPart p) => Semigroup1 (GenValue p)

instance AssocPair (GenProperty v an) where
  type Key (GenProperty v an) = Symbol () -- Don't want annotations to distinguish keys.
  type Value (GenProperty v an) = v an

  getKey = remAnns . propertyKey
  getValue = propertyValue

instance (Semigroup an) => SemAssocPair (GenProperty v an) where
  aappend1 vapp (Property xAnn xKey xVal) (Property yAnn yKey yVal)
    = Property (xAnn S.<> yAnn) (xKey `eappend` yKey) (xVal `vapp` yVal)

instance (Functor p, Foldable p, Traversable p) => SAnn (GenValue p) where
  mapSAnn f (Value ann parts) = Value (f ann) parts

instance (Functor p, Foldable p, Traversable p) => Ann (GenValue p) where
  getAnn = valueAnn

instance (Functor v, Foldable v, Traversable v) => Ann (GenRecord v) where
  getAnn = recordAnn

instance (Functor v, Foldable v, Traversable v) => Ann (GenProperty v) where
  getAnn = propertyAnn

instance (Printable p) => Printable (GenValue p) where
  aprintRec sub (Value _ parts) = pintercal " | " $ map sub parts

instance (Printable p) => FwdPrintable (GenValue p) where
  afprintRec = id

instance (FwdPrintable v, GenPropVal v) => Printable (GenRecord v) where
  aprintRec sub record = sub (head record) M.<> propsPrinted
    where propsPrinted = "[" M.<> pintercal ", " propPrinteds M.<> "]"
          propPrinteds = map sub $ properties record

instance (FwdPrintable v, GenPropVal v) => Printable (GenProperty v) where
  aprintRec sub (Property _ key val)
    | doesPrint val = pimp (keyPrinted M.<> ": ") M.<> valPrinted
    | otherwise = keyPrinted
    where keyPrinted = sub key
          valPrinted = afprintRec sub val

instance (Show an, Printable p, Show (p an)) => Summary (GenValue p an) where
  summaryRec = pprintSummaryRec

instance (Show an, FwdPrintable v, GenPropVal v, Show (v an)) => Summary (GenRecord v an) where
  summaryRec = pprintSummaryRec

instance (Show an, FwdPrintable v, GenPropVal v, Show (v an)) => Summary (GenProperty v an) where
  summaryRec = pprintSummaryRec

-- | Creates a value with just the given part.
singletonValue :: (Ann p) => p an -> GenValue p an
singletonValue part = Value (getAnn part) [part]

-- | Whether the value has any parts.
isEmpty :: GenValue p an -> Bool
isEmpty (Value _ parts) = null parts

-- | All primitives in the value.
primParts :: (GenPart p) => GenValue p an -> [Prim an]
primParts = mapMaybe partToPrim . valueParts

-- | Maps the parts.
mapParts :: (p1 an -> p2 an) -> GenValue p1 an -> GenValue p2 an
mapParts f (Value ann parts) = Value ann $ map f parts

-- | Maps then folds the parts.
foldMapParts :: (Monoid r) => (p an -> r) -> GenValue p an -> r
foldMapParts f (Value _ parts) = foldMap f parts

-- | Traverses the parts.
traverseParts :: (Applicative w)
              => (p1 an -> w (p2 an))
              -> GenValue p1 an
              -> w (GenValue p2 an)
traverseParts f (Value ann parts) = Value ann <$> traverse f parts

-- | Gets the record in the value with the given head.
recWithHead :: (GenPart p)
            => FSymbol an1
            -> GenValue p an2
            -> Maybe (PartRecord p an2)
recWithHead head' (Value _ parts)
  = find (recHasHead head') $ mapMaybe partToRec parts

-- | Gets the record in the value with the given head.
-- Returns an error if the value isn't present
forceRecWithHead :: (GenPart p)
                 => FSymbol an1
                 -> GenValue p an2
                 -> PartRecord p an2
forceRecWithHead head' value
  = case recWithHead head' value of
         Nothing
           -> error
            $ "Value doesn't contain record with head: "
           ++ show (remAnns head')
         Just record -> record

-- | Removes the record in the value with the given head.
-- If this record doesn't exist, does nothing.
deleteRecWithHead :: (GenPart p)
                  => FSymbol an
                  -> GenValue p an
                  -> GenValue p an
deleteRecWithHead head' (Value ann parts)
  = Value ann $ filter (partHasHead head') parts

cmergeParts :: (GenPart p, Semigroup an)
            => [[p an]]
            -> [p an]
cmergeParts = foldr mergeParts []

mergeParts :: (GenPart p, Semigroup an)
           => [p an]
           -> [p an]
           -> [p an]
mergeParts = foldr mergeAddPart

mergeAddRecord :: (GenPart p, Semigroup an)
               => PartRecord p an
               -> [p an]
               -> [p an]
mergeAddRecord record [] = [recToPart Proxy record]
mergeAddRecord record (x : xs)
  = case tryMergeRecordWithPart record x of
         Failure () -> x : mergeAddRecord record xs
         Success newRecord -> recToPart Proxy newRecord : xs

tryMergeRecordWithPart :: (GenPart p, Semigroup an)
                       => PartRecord p an
                       -> p an
                       -> UResult (PartRecord p an)
tryMergeRecordWithPart xRecord y
  = case partToRec y of
         Nothing -> Failure ()
         Just yRecord -> tryMergeRecords xRecord yRecord

tryMergeRecords :: (Semigroup1 v, Semigroup an)
                => GenRecord v an
                -> GenRecord v an
                -> UResult (GenRecord v an)
tryMergeRecords (Record xAnn xHead xProps) (Record yAnn yHead yProps)
  | xHead /@= yHead = Failure ()
  | otherwise = Success Record
  { recordAnn = xAnn S.<> yAnn
  , head = xHead `eappend` yHead
  , properties = assocUnionBy sappend1 xProps yProps
  }

-- | Whether the part's a record with the given head.
partHasHead :: (GenPart p) => FSymbol an -> p an -> Bool
partHasHead head' part
  = case partToRec part of
         Just record -> recHasHead head' record
         Nothing -> False

-- | Does the record have the given head?
recHasHead :: FSymbol an1 -> GenRecord v an2 -> Bool
recHasHead head' record = head' =@= head record

-- | Does the record have the given property key?
recHasProp :: Symbol an1 -> GenRecord v an2 -> Bool
recHasProp key = assocMember (remAnns key) . properties

-- | Gets the property with the given key.
-- Returns 'Nothing' if the property doesn't exist.
lookupProp :: Symbol an1 -> GenRecord v an2 -> Maybe (v an2)
lookupProp key = glookup (remAnns key) . properties

-- | Gets the property with the given key.
-- Raises an error if the property doesn't exist.
forceLookupProp :: Symbol an1 -> GenRecord v an2 -> v an2
forceLookupProp key record
  = case lookupProp key record of
         Nothing
           -> error
            $ "No property with the given key: "
           ++ show (remAnns key)
         Just x -> x

-- | Maps the record's property values.
mapPropVals :: (v an -> v2 an) -> GenRecord v an -> GenRecord v2 an
mapPropVals f record
  = Record
  { recordAnn = recordAnn record
  , head = head record
  , properties = map (mapPropVal f) $ properties record
  }

-- | Maps then folds the record's property values.
foldMapPropVals :: (Monoid r) => (v an -> r) -> GenRecord v an -> r
foldMapPropVals f = foldMap (f . propertyValue) . properties

-- | Traverses the record's property values.
traversePropVals :: (Applicative w) => (v an -> w (v2 an)) -> GenRecord v an -> w (GenRecord v2 an)
traversePropVals f record
    = Record (recordAnn record) (head record)
  <$> traverse (traversePropVal f) (properties record)

mapPropVal :: (v an -> v2 an) -> GenProperty v an -> GenProperty v2 an
mapPropVal f (Property ann key val) = Property ann key $ f val

traversePropVal :: (Applicative w) => (v an -> w (v2 an)) -> GenProperty v an -> w (GenProperty v2 an)
traversePropVal f (Property ann key val) = Property ann key <$> f val

-- | Uses the old parts to taint the value's annotation if necessary --
-- specifically, the value is tainted if the number of parts changes.
reconValue :: (GenPart p, TaintAnn an)
           => an
           -> [p an]
           -> [GenValue p an]
           -> GenValue p an
reconValue ann oldParts = reValue ann oldParts . cmergeParts . map valueParts

-- | Uses the old parts to taint the value's annotation if necessary --
-- specifically, the value is tainted if the number of parts changes.
reValue :: (TaintAnn an) => an -> [p an] -> [p an] -> GenValue p an
reValue ann oldParts newParts = Value ann' newParts
  where ann'
          | length oldParts /= length newParts = taint ann
          | otherwise = ann

-- | Uses the old properties to taint the record's annotation if
-- necessary -- specifically, the value is tainted if the number of
-- properties changes.
reRecord :: (TaintAnn an)
         => an
         -> FSymbol an
         -> [GenProperty v an]
         -> [GenProperty v an]
         -> GenRecord v an
reRecord ann head' oldProps newProps = Record ann' head' newProps
  where ann'
          | length oldProps /= length newProps = taint ann
          | otherwise = ann

-- | Uses the old property value to taint the property's annotation if
-- necessary -- specifically, the property is tainted if the old value
-- has a print and the new one doesn't, or vice versa.
reProperty :: (GenPropVal v, TaintAnn an)
           => an
           -> Symbol an
           -> v an
           -> v an
           -> GenProperty v an
reProperty ann key old new = Property ann' key new
  where ann'
          | doesPrint old /= doesPrint new = taint ann
          | otherwise = ann