{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | The versioned binary surface: encode a triangulation to bytes and read it
-- back. Decoding refuses a payload whose format version or coordinate encoding this
-- build does not own, rather than reinterpreting it.
module Moonlight.Triangulation.Serialization
  ( SerializationError (..)
  , serializationVersion
  , encodeTriangulation
  , decodeTriangulation
  ) where

import Control.Monad (replicateM, unless, when)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Except (ExceptT, runExceptT, throwE)
import Data.Binary (Binary (..))
import Data.Binary.Get
  ( Get
  , getDoublebe
  , getWord16be
  , getWord32be
  , getWord64be
  , getWord8
  , runGetOrFail
  )
import Data.Binary.Put
  ( Put
  , putDoublebe
  , putWord16be
  , putWord32be
  , putWord64be
  , putWord8
  , runPut
  )
import qualified Data.ByteString.Lazy as BL
import Data.Foldable (traverse_)
import qualified Data.IntSet as IntSet
import Data.Int (Int64)
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.Map.Strict as Map
import Data.Proxy (Proxy (..))
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import Data.Word (Word16, Word64, Word8)
import Moonlight.Triangulation.Internal.BoxedPaged (boxedFromVector, boxedToVector)
import Moonlight.Triangulation.Internal.Paged (fromLocalVector, fromVector, toVector)
import Moonlight.Triangulation.Internal.PointIndex (buildPointIndex)
import Moonlight.Triangulation.Handles.HandleDefs
import Moonlight.Triangulation.Internal.Representation
import Moonlight.Triangulation.Internal.Types
import Moonlight.Triangulation.Math (mkQueryPoint)
import Moonlight.Triangulation.Validation (validateTriangulation)

instance Binary (Point) where
  put (Point x y) = putDoublebe x >> putDoublebe y
  get = Point <$> getDoublebe <*> getDoublebe

instance Binary VertexId where
  put (VertexId value) = putWord32be value
  get = VertexId <$> getWord32be

instance Binary FaceId where
  put (FaceId value) = putWord32be value
  get = FaceId <$> getWord32be

instance Binary DirectedEdgeId where
  put (DirectedEdgeId value) = putWord32be value
  get = DirectedEdgeId <$> getWord32be

instance Binary UndirectedEdgeId where
  put (UndirectedEdgeId value) = putWord32be value
  get = UndirectedEdgeId <$> getWord32be

-- | Every way serialization refuses, each naming its witness.
data SerializationError
  = BinaryDecodeFailure !Int64 !String
  | TrailingBytes !Int64
  | InvalidFormatMagic !Word64
  | UnsupportedFormatVersion !Word16
  | ConstraintModeTagMismatch !Word8 !Word8
  | CoordinateEncodingTagMismatch !Word8 !Word8
  | EncodedCountExceedsInt !Word64
  | SerializedCoordinateLengthMismatch {-# UNPACK #-} !Int {-# UNPACK #-} !Int
  | SerializedVertexPayloadLengthMismatch {-# UNPACK #-} !Int {-# UNPACK #-} !Int
  | SerializedVertexOutgoingLengthMismatch {-# UNPACK #-} !Int {-# UNPACK #-} !Int
  | SerializedDirectedEdgeCountOdd {-# UNPACK #-} !Int
  | SerializedNextLengthMismatch {-# UNPACK #-} !Int {-# UNPACK #-} !Int
  | SerializedPreviousLengthMismatch {-# UNPACK #-} !Int {-# UNPACK #-} !Int
  | SerializedFaceReferenceLengthMismatch {-# UNPACK #-} !Int {-# UNPACK #-} !Int
  | SerializedDirectedPayloadLengthMismatch {-# UNPACK #-} !Int {-# UNPACK #-} !Int
  | SerializedUndirectedPayloadLengthMismatch {-# UNPACK #-} !Int {-# UNPACK #-} !Int
  | SerializedConstraintLengthMismatch {-# UNPACK #-} !Int {-# UNPACK #-} !Int
  | SerializedMissingOuterFace
  | SerializedFacePayloadLengthMismatch {-# UNPACK #-} !Int {-# UNPACK #-} !Int
  | InvalidSerializedPoint {-# UNPACK #-} !Int !PointValidationError
  | NonCanonicalSerializedConstraintFlag !UndirectedEdgeId !Word8
  | SerializedConstraintCountMismatch {-# UNPACK #-} !Int {-# UNPACK #-} !Int
  | UnconstrainedSerializedConstraints {-# UNPACK #-} !Int
  | DuplicateSerializedCoordinates {-# UNPACK #-} !Int {-# UNPACK #-} !Int
  | DecodedInvariantViolations !(NonEmpty InvariantViolation)
  deriving stock (Eq, Show)

type Decoder = ExceptT SerializationError Get

-- | The envelope version this module writes.
serializationVersion :: Word16
serializationVersion = 4

formatMagic :: Word64
formatMagic = 0x5350414445485307 -- "SPADEHS" + canonical geometry-owned format family

binary64EncodingTag :: Word8
binary64EncodingTag = 2

-- | Write the versioned binary envelope.
encodeTriangulation
  :: forall mode vertex directed undirected face. (KnownConstraintMode mode, Binary vertex, Binary directed, Binary undirected, Binary face)
  => Triangulation mode vertex directed undirected face
  -> BL.ByteString
encodeTriangulation triangulation = runPut $ do
  putWord64be formatMagic
  putWord16be serializationVersion
  putWord8 (modeTag (constraintModeValue (modeProxy triangulation)))
  putWord8 binary64EncodingTag
  let ElementDefaults directedDefault undirectedDefault faceDefault = triElementDefaults triangulation
  put directedDefault
  put undirectedDefault
  put faceDefault
  -- Geometry and payloads are independent components. Persist the authoritative
  -- coordinate pages rather than attempting to recover them from annotations.
  putUVector putDoublebe (toVector (triPointX triangulation))
  putUVector putDoublebe (toVector (triPointY triangulation))
  putBoxedVector put (boxedToVector (triVertexData triangulation))
  putUVector putWord32be (toVector (triVertexOut triangulation))
  -- The wire format stores the four topology planes separately; the interleaved
  -- arena is a resident layout, not a serialization concern.
  let topology = toVector (triHalfTopology triangulation)
      plane field = U.generate (U.length topology `quot` 4) (\edge -> topology U.! (4 * edge + field))
  putUVector putWord32be (plane 0)
  putUVector putWord32be (plane 1)
  putUVector putWord32be (plane 2)
  putUVector putWord32be (plane 3)
  putBoxedVector put (boxedToVector (triDirectedData triangulation))
  putBoxedVector put (boxedToVector (triUndirectedData triangulation))
  putUVector putWord32be (toVector (triFaceEdge triangulation))
  putBoxedVector put (boxedToVector (triFaceData triangulation))
  putUVector putWord8 (toVector (triConstraint triangulation))
  putCount (fromIntegral (triConstraintCount triangulation))

-- | Decode one exact, versioned finite DCEL. Coordinate uniqueness and the
-- complete topology, geometry, and Delaunay/CDT invariants are checked before
-- the opaque value is returned.
decodeTriangulation
  :: forall mode vertex directed undirected face.
     ( KnownConstraintMode mode
     , Binary vertex
     , Binary directed
     , Binary undirected
     , Binary face
     )
  => BL.ByteString
  -> Either SerializationError (Triangulation mode vertex directed undirected face)
decodeTriangulation bytes =
  case runGetOrFail (runExceptT getTriangulation) bytes of
    Left (_, offset, message) -> Left (BinaryDecodeFailure offset message)
    Right (_, _, Left failure) -> Left failure
    Right (remaining, _, Right triangulation)
      | not (BL.null remaining) -> Left (TrailingBytes (BL.length remaining))
      | otherwise ->
          case validateTriangulation triangulation of
            [] -> Right triangulation
            firstViolation : remainingViolations ->
              Left (DecodedInvariantViolations (firstViolation :| remainingViolations))
 where
  getTriangulation :: Decoder (Triangulation mode vertex directed undirected face)
  getTriangulation = do
    magic <- lift getWord64be
    unless (magic == formatMagic) (throwE (InvalidFormatMagic magic))
    version <- lift getWord16be
    unless (version == serializationVersion) (throwE (UnsupportedFormatVersion version))
    encodedMode <- lift getWord8
    let expectedMode = modeTag (constraintModeValue (Proxy :: Proxy mode))
    unless (encodedMode == expectedMode) (throwE (ConstraintModeTagMismatch expectedMode encodedMode))
    encodedScalar <- lift getWord8
    let expectedScalar = binary64EncodingTag
    unless (encodedScalar == expectedScalar) (throwE (CoordinateEncodingTagMismatch expectedScalar encodedScalar))
    defaults <- ElementDefaults <$> lift get <*> lift get <*> lift get
    pointXs <- getUVector (lift getDoublebe)
    pointYs <- getUVector (lift getDoublebe)
    vertexDataVector <- getBoxedVector (lift get)
    vertexOut <- getUVector (lift getWord32be)
    halfOrigin <- getUVector (lift getWord32be)
    halfNext <- getUVector (lift getWord32be)
    halfPrev <- getUVector (lift getWord32be)
    halfFace <- getUVector (lift getWord32be)
    directedDataVector <- getBoxedVector (lift get)
    undirectedDataVector <- getBoxedVector (lift get)
    faceEdge <- getUVector (lift getWord32be)
    faceDataVector <- getBoxedVector (lift get)
    constraints <- getUVector (lift getWord8)
    cachedConstraintCount <- getCount

    let vertexCount = U.length pointXs
        pointYCount = U.length pointYs
        vertexPayloadCount = V.length vertexDataVector
        vertexOutgoingCount = U.length vertexOut
        halfCount = U.length halfOrigin
        nextCount = U.length halfNext
        previousCount = U.length halfPrev
        faceReferenceCount = U.length halfFace
        edgeCount = halfCount `quot` 2
        faceCount = U.length faceEdge
        directedPayloadCount = V.length directedDataVector
        undirectedPayloadCount = V.length undirectedDataVector
        constraintCount = U.length constraints
        facePayloadCount = V.length faceDataVector
        points = zipWith Point (U.toList pointXs) (U.toList pointYs)
    unless (pointYCount == vertexCount) (throwE (SerializedCoordinateLengthMismatch vertexCount pointYCount))
    unless (vertexPayloadCount == vertexCount) (throwE (SerializedVertexPayloadLengthMismatch vertexPayloadCount vertexCount))
    unless (vertexOutgoingCount == vertexCount) (throwE (SerializedVertexOutgoingLengthMismatch vertexOutgoingCount vertexCount))
    unless (even halfCount) (throwE (SerializedDirectedEdgeCountOdd halfCount))
    unless (nextCount == halfCount) (throwE (SerializedNextLengthMismatch nextCount halfCount))
    unless (previousCount == halfCount) (throwE (SerializedPreviousLengthMismatch previousCount halfCount))
    unless (faceReferenceCount == halfCount) (throwE (SerializedFaceReferenceLengthMismatch faceReferenceCount halfCount))
    unless (directedPayloadCount == halfCount) (throwE (SerializedDirectedPayloadLengthMismatch directedPayloadCount halfCount))
    unless (undirectedPayloadCount == edgeCount) (throwE (SerializedUndirectedPayloadLengthMismatch undirectedPayloadCount edgeCount))
    unless (constraintCount == edgeCount) (throwE (SerializedConstraintLengthMismatch constraintCount edgeCount))
    unless (faceCount >= 1) (throwE SerializedMissingOuterFace)
    unless (facePayloadCount == faceCount) (throwE (SerializedFacePayloadLengthMismatch facePayloadCount faceCount))
    traverse_ (uncurry validateStoredPoint) (zip [0 ..] points)

    case U.ifoldr (\index flag found -> if flag /= 0 && flag /= 1 then Just (index, flag) else found) Nothing constraints of
      Nothing -> pure ()
      Just (index, flag) ->
        throwE
          ( NonCanonicalSerializedConstraintFlag
              (UndirectedEdgeId (fromIntegral index))
              flag
          )
    let actualConstraintCount = U.foldl' (\count flag -> if flag == 1 then count + 1 else count) 0 constraints
    unless (cachedConstraintCount == actualConstraintCount) (throwE (SerializedConstraintCountMismatch cachedConstraintCount actualConstraintCount))
    when (expectedMode == 0 && actualConstraintCount /= 0) (throwE (UnconstrainedSerializedConstraints actualConstraintCount))

    let distinctPoints = Map.fromList (map (\point -> (point, ())) points)
        distinctPointCount = Map.size distinctPoints
    unless (distinctPointCount == vertexCount) (throwE (DuplicateSerializedCoordinates vertexCount distinctPointCount))

    let pointXStore = fromLocalVector 0 pointXs
        pointYStore = fromLocalVector 0 pointYs
        topologyStore =
          fromVector maxBound $
            U.generate (4 * halfCount) $ \slot ->
              let (edge, field) = slot `quotRem` 4
               in case field of
                    0 -> halfOrigin U.! edge
                    1 -> halfNext U.! edge
                    2 -> halfPrev U.! edge
                    _ -> halfFace U.! edge
        constraintEdgeIndex =
          U.ifoldl'
            (\edges index flag ->
               if flag == 1 then IntSet.insert index edges else edges
            )
            IntSet.empty
            constraints
    pure
      Triangulation
        { triPointX = pointXStore
        , triPointY = pointYStore
        , triPointIndex = buildPointIndex pointXStore pointYStore
        , triVertexOut = fromLocalVector maxBound vertexOut
        , triVertexData = boxedFromVector Nothing vertexDataVector
        , triHalfTopology = topologyStore
        , triDirectedData = boxedFromVector (Just (defaultDirectedEdgeData defaults)) directedDataVector
        , triUndirectedData = boxedFromVector (Just (defaultUndirectedEdgeData defaults)) undirectedDataVector
        , triFaceEdge = fromLocalVector maxBound faceEdge
        , triFaceData = boxedFromVector (Just (defaultFaceData defaults)) faceDataVector
        , triConstraint = fromVector 0 constraints
        , triConstraintCount = cachedConstraintCount
        , triConstraintEdges = constraintEdgeIndex
        , triElementDefaults = defaults
        }

validateStoredPoint :: Int -> Point -> Decoder ()
validateStoredPoint index point =
  case mkQueryPoint point of
    Left failure -> throwE (InvalidSerializedPoint index failure)
    Right _ -> pure ()

modeProxy :: Triangulation mode vertex directed undirected face -> Proxy mode
modeProxy _ = Proxy

modeTag :: ConstraintMode -> Word8
modeTag Unconstrained = 0
modeTag Constrained = 1

putCount :: Word64 -> Put
putCount = putWord64be

getCount :: Decoder Int
getCount = do
  raw <- lift getWord64be
  when (raw > fromIntegral (maxBound :: Int)) (throwE (EncodedCountExceedsInt raw))
  pure (fromIntegral raw)

putUVector :: U.Unbox a => (a -> Put) -> U.Vector a -> Put
putUVector putElement values = do
  putCount (fromIntegral (U.length values))
  U.mapM_ putElement values

getUVector :: U.Unbox a => Decoder a -> Decoder (U.Vector a)
getUVector getElement = do
  count <- getCount
  U.fromList <$> replicateM count getElement

putBoxedVector :: (a -> Put) -> V.Vector a -> Put
putBoxedVector putElement values = do
  putCount (fromIntegral (V.length values))
  V.mapM_ putElement values

getBoxedVector :: Decoder a -> Decoder (V.Vector a)
getBoxedVector getElement = do
  count <- getCount
  V.fromList <$> replicateM count getElement