content: {-# LANGUAGE AllowAmbiguousTypes #-}
  
  module Telescope.Asdf.Class where
  
  import Data.List.NonEmpty (NonEmpty (..))
  import Data.List.NonEmpty qualified as NE
  import Data.Massiv.Array (Array, Prim)
  import Data.Massiv.Array qualified as M
  import Data.Scientific (fromFloatDigits, toRealFloat)
  import Data.Text (Text, pack, unpack)
  import Data.Time.Clock (UTCTime)
  import Data.Time.Format.ISO8601
  import Data.Time.LocalTime (LocalTime)
  import Effectful
  import Effectful.Fail
  import GHC.Generics
  import GHC.Int
  import Telescope.Asdf.Encoding.File (BlockSource (..))
  import Telescope.Asdf.NDArray
  import Telescope.Asdf.Node
  import Telescope.Data.Array
  import Telescope.Data.Axes
  import Telescope.Data.Binary
  import Telescope.Data.Parser
  
  
  {- | Convert a type to an Asdf 'Value' or 'Node'. The generic instance will encode to an 'Object' with field names matching record selectors
  
  > data Example = Example
  >   { name :: Text
  >   , age :: Int
  >   , tags :: [Text]
  >   }
  >   deriving (Generic, ToAsdf)
  >
  > instance ToAsdf Example where
  >   schema _ = "tag:example.org/schemas/example-1.0.0"
  -}
  class ToAsdf a where
    -- | Specify how an object encodes to a 'Value'
    --
    -- > instance User ToAsdf where
    -- >   toValue user =
    -- >     Object
    -- >       [ ("name", toNode user.name)
    -- >       , ("age", toNode user.age)
    -- >       ]
    toValue :: a -> Value
    default toValue :: (Generic a, GToObject (Rep a)) => a -> Value
    toValue a = Object $ gToObject (from a)
  
  
    -- | Specify the schema for a type
    --
    -- > instance ToAsdf Unit where
    -- >   schema _ = "!unit/unit-1.0.0"
    schema :: a -> SchemaTag
    default schema :: a -> SchemaTag
    schema _ = mempty
  
  
    -- | Specify that this node be saved as an anchor
    --
    -- > instance ToAsdf Config where
    -- >   anchor _ = Just "globalConfig"
    anchor :: a -> Maybe Anchor
    default anchor :: a -> Maybe Anchor
    anchor _ = Nothing
  
  
    -- | Manually control all aspects of how this is converted to a 'Node'
    toNode :: a -> Node
    default toNode :: a -> Node
    toNode a = Node (schema a) (anchor a) $ toValue a
  
  
  {- | Parse an Asdf 'Value' or 'Node' into a type. The generic instance will decode an 'Object' with field names matching record selectors
  
  > data Example = Example
  >   { name :: Text
  >   , age :: Int
  >   , tags :: [Text]
  >   }
  >   deriving (Generic, FromAsdf)
  -}
  class FromAsdf a where
    -- | Specify how a type is parsed from a 'Value'
    --
    -- > instance FromAsdf Integer where
    -- >   parseValue = \case
    -- >     Integer n -> pure $ fromIntegral n
    -- >     node -> expected "Integer" node
    parseValue :: (Parser :> es) => Value -> Eff es a
    default parseValue :: (Generic a, GParseObject (Rep a), Parser :> es) => Value -> Eff es a
    parseValue (Object o) = to <$> gParseObject o
    parseValue val = expected "Object" val
  
  
    parseNode :: (Parser :> es) => Node -> Eff es a
    parseNode (Node _ _ v) = parseValue v
  
  
  instance ToAsdf Int where
    toValue n = toValue (fromIntegral @Int @Int64 n)
  instance FromAsdf Int where
    parseValue = fmap (fromIntegral @Int64 @Int) <$> parseValue
  
  
  instance ToAsdf Int8 where
    toValue n = Integer $ fromIntegral n
  instance FromAsdf Int8 where
    parseValue = parseInteger
  
  
  instance ToAsdf Int16 where
    toValue n = Integer $ fromIntegral n
  instance FromAsdf Int16 where
    parseValue = parseInteger
  
  
  instance ToAsdf Int32 where
    toValue n = Integer $ fromIntegral n
  instance FromAsdf Int32 where
    parseValue = parseInteger
  
  
  instance ToAsdf Int64 where
    toValue n = Integer $ fromIntegral n
  instance FromAsdf Int64 where
    parseValue = parseInteger
  
  
  instance ToAsdf Integer where
    toValue n = Integer $ fromIntegral n
  instance FromAsdf Integer where
    parseValue = parseInteger
  
  
  instance ToAsdf Double where
    toValue n = Number $ fromFloatDigits n
  instance FromAsdf Double where
    parseValue = \case
      Number n -> pure $ toRealFloat n
      node -> expected "Double" node
  
  
  instance ToAsdf Float where
    toValue n = Number $ fromFloatDigits n
  instance FromAsdf Float where
    parseValue = \case
      Number n -> pure $ toRealFloat n
      node -> expected "Float" node
  
  
  parseInteger :: (Integral a, Parser :> es) => Value -> Eff es a
  parseInteger = \case
    Integer n -> pure $ fromIntegral n
    node -> expected "Integer" node
  
  
  instance {-# OVERLAPPABLE #-} (FromAsdf a) => FromAsdf [a] where
    parseValue = \case
      Array ns -> mapM (parseNode @a) ns
      node -> expected "Array" node
  instance {-# OVERLAPPABLE #-} (ToAsdf a) => ToAsdf [a] where
    toValue as = Array $ fmap toNode as
  instance {-# OVERLAPPABLE #-} (FromAsdf a) => FromAsdf (NonEmpty a) where
    parseValue val = do
      as <- parseValue @[a] val
      case as of
        [] -> expected "NonEmpty List" val
        (a : rest) -> pure (a :| rest)
  instance {-# OVERLAPPABLE #-} (ToAsdf a) => ToAsdf (NonEmpty a) where
    toValue as = toValue $ NE.toList as
  
  
  instance {-# OVERLAPS #-} ToAsdf Object where
    toValue = Object
  instance {-# OVERLAPS #-} FromAsdf Object where
    parseValue = \case
      Object o -> pure o
      node -> expected "Object" node
  
  
  instance ToAsdf () where
    toValue () = Null
  instance FromAsdf () where
    parseValue _ = pure ()
  
  
  instance (ToAsdf a, ToAsdf b) => ToAsdf (a, b) where
    toValue (a, b) = Array [toNode a, toNode b]
  instance (FromAsdf a, FromAsdf b) => FromAsdf (a, b) where
    parseValue = \case
      Array [n1, n2] -> do
        a <- parseValue n1.value
        b <- parseValue n2.value
        pure (a, b)
      node -> expected "[a, b]" node
  
  
  instance (ToAsdf a, ToAsdf b, ToAsdf c) => ToAsdf (a, b, c) where
    toValue (a, b, c) = Array [toNode a, toNode b, toNode c]
  instance (FromAsdf a, FromAsdf b, FromAsdf c) => FromAsdf (a, b, c) where
    parseValue = \case
      Array [na, nb, nc] -> do
        a <- parseValue na.value
        b <- parseValue nb.value
        c <- parseValue nc.value
        pure (a, b, c)
      node -> expected "[a, b, c]" node
  
  
  instance (ToAsdf a, ToAsdf b, ToAsdf c, ToAsdf d) => ToAsdf (a, b, c, d) where
    toValue (a, b, c, d) = Array [toNode a, toNode b, toNode c, toNode d]
  instance (FromAsdf a, FromAsdf b, FromAsdf c, FromAsdf d) => FromAsdf (a, b, c, d) where
    parseValue = \case
      Array [na, nb, nc, nd] -> do
        a <- parseValue na.value
        b <- parseValue nb.value
        c <- parseValue nc.value
        d <- parseValue nd.value
        pure (a, b, c, d)
      node -> expected "[a, b, c, d]" node
  
  
  instance (ToAsdf a, ToAsdf b, ToAsdf c, ToAsdf d, ToAsdf e) => ToAsdf (a, b, c, d, e) where
    toValue (a, b, c, d, e) = Array [toNode a, toNode b, toNode c, toNode d, toNode e]
  instance (FromAsdf a, FromAsdf b, FromAsdf c, FromAsdf d, FromAsdf e) => FromAsdf (a, b, c, d, e) where
    parseValue = \case
      Array [na, nb, nc, nd] -> do
        a <- parseValue na.value
        b <- parseValue nb.value
        c <- parseValue nc.value
        d <- parseValue nd.value
        e <- parseValue nd.value
        pure (a, b, c, d, e)
      node -> expected "[a, b, c, d, e]" node
  
  
  instance (ToAsdf a, ToAsdf b, ToAsdf c, ToAsdf d, ToAsdf e, ToAsdf f) => ToAsdf (a, b, c, d, e, f) where
    toValue (a, b, c, d, e, f) = Array [toNode a, toNode b, toNode c, toNode d, toNode e, toNode f]
  instance (FromAsdf a, FromAsdf b, FromAsdf c, FromAsdf d, FromAsdf e, FromAsdf f) => FromAsdf (a, b, c, d, e, f) where
    parseValue = \case
      Array [na, nb, nc, nd] -> do
        a <- parseValue na.value
        b <- parseValue nb.value
        c <- parseValue nc.value
        d <- parseValue nd.value
        e <- parseValue nd.value
        f <- parseValue nd.value
        pure (a, b, c, d, e, f)
      node -> expected "[a, b, c, d, e, f]" node
  
  
  -- they will always serialize to Array
  instance FromAsdf [Text] where
    parseValue = parseAnyList
  instance FromAsdf [Int] where
    parseValue = parseAnyList
  instance FromAsdf [Int8] where
    parseValue = parseAnyList
  instance FromAsdf [Int16] where
    parseValue = parseAnyList
  instance FromAsdf [Int32] where
    parseValue = parseAnyList
  instance FromAsdf [Int64] where
    parseValue = parseAnyList
  instance FromAsdf [Double] where
    parseValue = parseAnyList
  
  
  -- | Flexibly parse lists from either Array or NDArray
  parseAnyList :: (FromAsdf a, FromNDArray [a], Parser :> es) => Value -> Eff es [a]
  parseAnyList = \case
    Array ns -> mapM parseNode ns
    NDArray dat -> fromNDArray dat
    node -> expected "[Double]" node
  
  
  instance (FromAsdf a) => FromAsdf (Maybe a) where
    parseNode node =
      case node.value of
        Null -> pure Nothing
        _ -> Just <$> parseNode @a node
    parseValue = \case
      Null -> pure Nothing
      val -> Just <$> parseValue @a val
  instance (ToAsdf a) => ToAsdf (Maybe a) where
    schema = maybe mempty schema
    anchor = maybe Nothing anchor
    toValue Nothing = Null
    toValue (Just a) = toValue a
  
  
  instance (BinaryValue a, Prim a, AxesIndex ix) => FromAsdf (Array M.D ix a) where
    parseValue = \case
      NDArray a -> fromNDArray a
      node -> expected "NDArray" node
  instance (BinaryValue a, IsDataType a, Prim a, AxesIndex ix, PutArray ix) => ToAsdf (Array M.D ix a) where
    toValue as = NDArray $ ndArrayMassiv as
  
  
  instance ToAsdf Text where
    toValue = String
  instance FromAsdf Text where
    parseValue = \case
      String t -> pure t
      node -> expected "Text" node
  
  
  instance ToAsdf String where
    toValue = String . pack
  instance FromAsdf String where
    parseValue = \case
      String t -> pure $ unpack t
      node -> expected "String" node
  
  
  instance ToAsdf Bool where
    toValue = Bool
  instance FromAsdf Bool where
    parseValue = \case
      Bool b -> pure b
      node -> expected "Bool" node
  
  
  instance ToAsdf Value where
    toValue = id
  instance FromAsdf Value where
    parseValue = pure
  
  
  instance ToAsdf Node where
    schema n = n.schema
    anchor n = n.anchor
    toNode n = n
    toValue (Node _ _ val) = val
  instance FromAsdf Node where
    parseNode = pure
    parseValue val = pure $ Node mempty Nothing val
  
  
  instance ToAsdf Tree where
    toValue (Tree o) = Object o
  instance FromAsdf Tree where
    parseValue = \case
      Object o -> pure $ Tree o
      val -> expected "Object" val
  
  
  instance ToAsdf NDArrayData where
    toValue = NDArray
  instance FromAsdf NDArrayData where
    parseValue = \case
      NDArray nda -> pure nda
      node -> expected "NDArray" node
  
  
  instance ToAsdf DataType where
    toValue Float64 = "float64"
    toValue Float32 = "float32"
    toValue Int64 = "int64"
    toValue Int32 = "int32"
    toValue Int16 = "int16"
    toValue Int8 = "int8"
    toValue Bool8 = "bool8"
    toValue (Ucs4 n) = Array ["ucs4", fromValue $ Integer $ fromIntegral n]
  instance FromAsdf DataType where
    parseValue = \case
      String "float64" -> pure Float64
      String "float32" -> pure Float32
      String "int64" -> pure Int64
      String "int32" -> pure Int32
      String "int16" -> pure Int16
      String "int8" -> pure Int8
      String "bool8" -> pure Bool8
      Array ["ucs4", Node _ _ (Integer n)] -> pure $ Ucs4 $ fromIntegral n
      val -> expected "DataType" val
  
  
  instance ToAsdf ByteOrder where
    toValue = \case
      BigEndian -> "big"
      LittleEndian -> "little"
  instance FromAsdf ByteOrder where
    parseValue = \case
      String "big" -> pure BigEndian
      String "little" -> pure LittleEndian
      node -> expected "ByteOrder" node
  
  
  instance ToAsdf (Axes Row) where
    toValue (Axes as) = toValue as
  
  
  instance ToAsdf BlockSource where
    toValue (BlockSource s) = toValue s
  
  
  instance ToAsdf UTCTime where
    schema _ = "!time/time-1.1.0"
    toValue t = String $ pack $ iso8601Show t
  instance FromAsdf UTCTime where
    parseValue v = do
      ts <- parseValue @String v
      res <- runFail $ iso8601ParseM ts
      case res of
        Left e -> parseFail e
        Right a -> pure a
  instance ToAsdf LocalTime where
    schema _ = "!time/time-1.1.0"
    toValue t = String $ pack $ iso8601Show t
  instance FromAsdf LocalTime where
    parseValue v = do
      ts <- parseValue @String v
      res <- runFail $ iso8601ParseM ts
      case res of
        Left e -> parseFail e
        Right a -> pure a
  
  
  {- | Parse a key from an 'Object'
  
  > instance FromAsdf User where
  >   parseValue = \case
  >     Object o -> do
  >       name <- o .: "name"
  >       age <- o .: "age"
  >       pure $ User{name, age}
  -}
  (.:) :: (FromAsdf a, Parser :> es) => Object -> Key -> Eff es a
  o .: k = do
    case lookup k o of
      Nothing -> parseFail $ "key " ++ show k ++ " not found"
      Just node ->
        parseAt (Child k) $ parseNode node
  
  
  -- | Parse an optional key from an 'Object'
  (.:?) :: (FromAsdf a, Parser :> es) => Object -> Key -> Eff es (Maybe a)
  o .:? k = do
    case lookup k o of
      Nothing -> pure Nothing
      Just a -> do
        parseAt (Child k) $ parseNode a
  
  
  {- | Parse a child at the given array index
   -
  > instance FromAsdf Friends where
  >   parseValue = \case
  >     Array ns -> do
  >       best <- ns ! 0
  >       second <- ns ! 1
  >       other <- mapM parseNode ns
  >       pure $ Friends{best, second, other}
  -}
  (!) :: (FromAsdf a, Parser :> es) => [Node] -> Int -> Eff es a
  ns ! n = do
    if n >= length ns
      then parseFail $ "Index " ++ show n ++ " not found"
      else parseAt (Index n) $ parseNode $ ns !! n
  
  
  -- | Generically serialize records to an 'Object'
  class GToObject f where
    gToObject :: f p -> Object
  
  
  instance (GToObject f) => GToObject (M1 D c f) where
    gToObject (M1 f) = gToObject f
  
  
  instance (GToObject f) => GToObject (M1 C c f) where
    gToObject (M1 f) = gToObject f
  
  
  instance (GToObject f, GToObject g) => GToObject (f :*: g) where
    gToObject (f :*: g) = gToObject f <> gToObject g
  
  
  instance (GToNode f, Selector s) => GToObject (M1 S s f) where
    gToObject (M1 f) =
      let s = selName (undefined :: M1 S s f p)
       in [(pack s, gToNode f)]
  
  
  -- | Generically serialize record values to a 'Node'
  class GToNode f where
    gToNode :: f p -> Node
  
  
  instance {-# OVERLAPPABLE #-} (ToAsdf a) => GToNode (K1 R a) where
    gToNode (K1 a) = toNode a
  
  
  instance {-# OVERLAPPING #-} (ToAsdf a) => GToNode (K1 R (Maybe a)) where
    gToNode (K1 a) = toNode a
  
  
  -- | Generically parse 'Object's into records
  class GParseObject f where
    gParseObject :: (Parser :> es) => Object -> Eff es (f p)
  
  
  instance (GParseObject f) => GParseObject (M1 D c f) where
    gParseObject o = M1 <$> gParseObject o
  
  
  instance (GParseObject f) => GParseObject (M1 C c f) where
    gParseObject o = M1 <$> gParseObject o
  
  
  instance (GParseObject f, GParseObject g) => GParseObject (f :*: g) where
    gParseObject o = do
      f <- gParseObject o
      g <- gParseObject o
      pure $ f :*: g
  
  
  instance (GParseKey f, Selector s) => GParseObject (M1 S s f) where
    gParseObject o = do
      let k = pack $ selName (undefined :: M1 S s f p)
      M1 <$> gParseKey o k
  
  
  -- | Generically parse a key from an 'Object' into a record value
  class GParseKey f where
    gParseKey :: (Parser :> es) => Object -> Key -> Eff es (f p)
  
  
  instance {-# OVERLAPPABLE #-} (FromAsdf a) => GParseKey (K1 R a) where
    gParseKey o k = K1 <$> o .: k
  
  
  instance {-# OVERLAPPABLE #-} (FromAsdf a) => GParseKey (K1 R (Maybe a)) where
    gParseKey o k = K1 <$> o .:? k
  