{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}

module Database.PostgreSQL.Opium.FromField
  ( -- * Decoding data from @libpq@
    FromField (..)
  , fromField
    -- * Utility types
  , RawField (..)
  ) where

import Data.Attoparsec.ByteString (Parser)
import Data.Bits (Bits (..))
import Data.ByteString (ByteString)
import Data.Functor (($>))
import Data.Int (Int16, Int32)
import Data.Proxy (Proxy (..))
import Data.Time
  ( Day (..)
  , DiffTime
  , LocalTime (..)
  , TimeOfDay
  , addDays
  , fromGregorian
  , picosecondsToDiffTime
  , timeToTimeOfDay
  )
import Data.Text (Text)
import Data.Word (Word16, Word32)
import Database.PostgreSQL.LibPQ (Oid)
import Unsafe.Coerce (unsafeCoerce)

import qualified Data.Attoparsec.ByteString as AP
import qualified Data.ByteString as BS
import qualified Data.Text as Text
import qualified Data.Text.Encoding as Encoding

import qualified Database.PostgreSQL.Opium.Oid as Oid

(\/) :: (a -> Bool) -> (a -> Bool) -> a -> Bool
p \/ q = \x -> p x || q x

fromField :: FromField a => ByteString -> Either String a
fromField =
  AP.parseOnly parseField

class FromField a where
  validOid :: Proxy a -> Oid -> Bool
  parseField :: Parser a

-- | See https://www.postgresql.org/docs/current/datatype-binary.html.
-- Accepts @bytea@.
instance FromField ByteString where
  validOid Proxy = Oid.bytea
  parseField = AP.takeByteString

-- | See https://www.postgresql.org/docs/current/datatype-character.html.
-- Accepts @text@, @character@ and @character varying@.
instance FromField Text where
  validOid Proxy = Oid.text \/ Oid.character \/ Oid.characterVarying
  parseField = Encoding.decodeUtf8 <$> AP.takeByteString

-- Accepts @text@, @character@ and @character varying@.
-- | See https://www.postgresql.org/docs/current/datatype-character.html.
instance FromField String where
  validOid Proxy = validOid @Text Proxy
  parseField = Text.unpack <$> parseField

-- | See https://www.postgresql.org/docs/current/datatype-character.html.
-- This instance accepts all character types but fails to decode fields that are not exactly one character.
instance FromField Char where
  validOid Proxy = validOid @Text Proxy
  parseField = do
    str <- parseField
    case str of
      [c] -> pure c
      _ -> fail "Char accepts single characters only"


readBigEndian :: (Bits a, Num a) => ByteString -> a
readBigEndian = BS.foldl' (\x b -> x `shiftL` 8 .|. fromIntegral b) 0

readInt :: Num a => ByteString -> Parser a
readInt bs = case BS.length bs of
  4 -> pure $ fromIntegral $ readBigEndian @Int32 bs
  8 -> pure $ fromIntegral $ readBigEndian @Int bs
  2 -> pure $ fromIntegral $ readBigEndian @Int16 bs
  _ -> fail "Wrong number bytes for integer"

readWord :: Num a => ByteString -> Parser a
readWord bs = case BS.length bs of
  4 -> pure $ fromIntegral $ readBigEndian @Word32 bs
  8 -> pure $ fromIntegral $ readBigEndian @Word bs
  2 -> pure $ fromIntegral $ readBigEndian @Word16 bs
  _ -> fail "Wrong number bytes for word"

-- | See https://www.postgresql.org/docs/current/datatype-numeric.html.
-- We assume that 'Int' has 64 bits. This is not guaranteed but reasonable enough.
instance FromField Int where
  validOid Proxy = Oid.smallint \/ Oid.integer \/ Oid.bigint
  parseField = readInt =<< AP.takeByteString

-- | See https://www.postgresql.org/docs/current/datatype-numeric.html.
instance FromField Integer where
  validOid Proxy = Oid.smallint \/ Oid.integer \/ Oid.bigint
  parseField = readInt =<< AP.takeByteString

instance FromField Word where
  validOid Proxy = Oid.smallint \/ Oid.integer \/ Oid.bigint
  parseField = readWord =<< AP.takeByteString

-- | See https://www.postgresql.org/docs/current/datatype-numeric.html.
-- Accepts only @real@ fields, not @double precision@.
instance FromField Float where
  validOid Proxy = Oid.real
  -- Afaict there's no cleaner (@base@) way to access the underlying bits.
  -- In C we'd do
  --
  -- union { float a; uint32_t b; } x;
  -- x.b = ...;
  -- return x.a;
  parseField = unsafeCoerce <$> readBigEndian @Word32 <$> AP.takeByteString

-- | See https://www.postgresql.org/docs/current/datatype-numeric.html.
-- Accepts only @double precision@ fields, not @real@.
instance FromField Double where
  validOid Proxy = Oid.doublePrecision
  parseField = unsafeCoerce <$> readBigEndian @Word <$> AP.takeByteString

boolParser :: Parser Bool
boolParser = AP.choice
  [ AP.word8 1 $> True
  , AP.word8 0 $> False
  ]

-- | See https://www.postgresql.org/docs/current/datatype-boolean.html.
instance FromField Bool where
  validOid Proxy = Oid.boolean
  parseField = boolParser

postgresEpoch :: Day
postgresEpoch = fromGregorian 2000 1 1

fromPostgresJulian :: Integer -> Day
fromPostgresJulian x = addDays x postgresEpoch

-- | See https://www.postgresql.org/docs/current/datatype-datetime.html.
-- Relevant as well: https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/datetime.c;h=267dfd37b2e8b9bc63797c69b9ca2e45e6bfde61;hb=HEAD#l267.
-- Note that Postgres uses the proleptic Gregorian calendar, whereas @Show Day@ and @fromGregorian@ use an astronomical calendar.
-- In short, Postgres treats 1 BC as a leap year and doesn't have a year zero.
-- This means that working with negative dates will be different in Postgres and your application code.
instance FromField Day where
  validOid Proxy = Oid.date
  parseField = fromPostgresJulian . fromIntegral <$> readBigEndian @Int32 <$> AP.takeByteString

-- | See https://www.postgresql.org/docs/current/datatype-datetime.html.
-- Binary format: https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/date.c;h=ae0f24de2c3c54eb6d0405cdb212597c2407238e;hb=HEAD#l1542.
-- Accepts @time@.
instance FromField DiffTime where
  validOid Proxy = Oid.time
  parseField = microsecondsToDiffTime . fromIntegral <$> readBigEndian @Int <$> AP.takeByteString
    where
      microsecondsToDiffTime :: Integer -> DiffTime
      microsecondsToDiffTime ms = picosecondsToDiffTime $ ms * 1000000

-- | See https://www.postgresql.org/docs/current/datatype-datetime.html.
-- Binary format: https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/date.c;h=ae0f24de2c3c54eb6d0405cdb212597c2407238e;hb=HEAD#l1542.
-- Accepts @time@.
instance FromField TimeOfDay where
  validOid Proxy = Oid.time
  parseField = timeToTimeOfDay <$> parseField @DiffTime

-- | See https://www.postgresql.org/docs/current/datatype-datetime.html.
-- | Accepts @timestamp@.
-- Note that Postgres uses the proleptic Gregorian calendar, whereas @Show Day@ and @fromGregorian@ use an astronomical calendar.
-- In short, Postgres treats 1 BC as a leap year and doesn't have a year zero.
-- This means that working with negative dates will be different in Postgres and your application code.
instance FromField LocalTime where
  validOid Proxy = Oid.timestamp
  parseField = fromPostgresTimestamp <$> readBigEndian @Int <$> AP.takeByteString
    where
      fromPostgresTimestamp :: Int -> LocalTime
      fromPostgresTimestamp ts =
        let
          (days, microseconds) = ts `divMod` (86400 * 1000000)
          day = fromPostgresJulian $ fromIntegral days
          time = picosecondsToDiffTime $ fromIntegral microseconds * 1000000
        in
          LocalTime day (timeToTimeOfDay time)

newtype RawField a = RawField a
  deriving (Eq, Show)

instance FromField a => FromField (RawField a) where
  validOid Proxy = const True
  parseField = RawField <$> parseField