content: module Telescope.Fits.Encoding where
  
  import Control.Monad.Catch (MonadThrow, throwM)
  import Data.ByteString (ByteString)
  import Data.ByteString qualified as BS
  import Data.Text (Text)
  import Data.Text.Encoding qualified as TE
  import Effectful
  import Effectful.Dispatch.Dynamic
  import Effectful.Error.Static
  import Effectful.State.Static.Local
  import Telescope.Data.Parser
  import Telescope.Fits.Checksum
  import Telescope.Fits.DataArray (DataArray (..), Dimensions (..), dataArray, dataSizeBytes)
  import Telescope.Fits.Encoding.MegaHeader qualified as MH
  import Telescope.Fits.Encoding.Render
  import Telescope.Fits.HDU
  import Telescope.Fits.Header
  import Text.Megaparsec (lookAhead)
  import Text.Megaparsec qualified as MP
  
  
  {- | Decode a FITS file read from a strict 'ByteString'
  
  >  decode =<< BS.readFile "samples/simple2x3.fits"
  -}
  decode :: forall m. (MonadThrow m) => ByteString -> m Fits
  decode = fitsParseThrow parseFits
  
  
  {- | Encode a FITS file to a strict 'ByteString'
  
  > BS.writeFile $ encode fits
  -}
  encode :: Fits -> ByteString
  encode f =
    let prim = encodePrimaryHDU f.primaryHDU
        exts = fmap encodeExtension f.extensions
     in mconcat $ prim : exts
  
  
  fitsParseThrow :: (MonadThrow m) => Eff '[State ByteString, Parser, Error ParseError] a -> ByteString -> m a
  fitsParseThrow parse inp =
    case runFitsParse parse inp of
      Left e -> throwM e
      Right a -> pure a
  
  
  runFitsParse :: Eff '[State ByteString, Parser, Error ParseError] a -> BS.ByteString -> Either ParseError a
  runFitsParse parse inp = do
    runParserPure $ evalState inp parse
  
  
  parseFits :: forall es. (State ByteString :> es, Parser :> es) => Eff es Fits
  parseFits = do
    p <- primary
    es <- parseAt (Child "extensions") $ extensions 1
    pure $ Fits p es
  
  
  primary :: (State ByteString :> es, Parser :> es) => Eff es DataHDU
  primary = do
    parseAt (Child "primary") $ do
      (dm, hd) <- runMega "Primary Header" $ do
        dm <- MP.lookAhead MH.parsePrimaryKeywords
        hd <- MH.parseHeader
        pure (dm, hd)
      darr <- mainData dm
      pure $ DataHDU hd darr
  
  
  extensions :: (State ByteString :> es, Parser :> es) => Int -> Eff es [Extension]
  extensions n = do
    inp <- get @ByteString
    -- recurse until you run out of input or error
    case inp of
      "" -> pure []
      _ -> do
        e <- parseAt (Index n) extension
        es <- extensions (n + 1)
        pure (e : es)
  
  
  extension :: (State ByteString :> es, Parser :> es) => Eff es Extension
  extension = do
    -- I don't understand how to use NonDet to fix this
    resImg <- runErrorNoCallStack @ParseError $ runParser image
    resTbl <- runErrorNoCallStack @ParseError $ runParser binTable
    case (resImg, resTbl) of
      (Right i, _) -> pure $ Image i
      (_, Right b) -> pure $ BinTable b
      -- (Left (_, FormatError ie), Left (_, FormatError be)) -> throwM $ InvalidHDU [ie, be]
      (Left _, Left (ParseFailure p e)) -> do
        send $ PathMod (<> p) $ parseFail e
  
  
  image :: (State ByteString :> es, Parser :> es) => Eff es DataHDU
  image = do
    (dm, hd) <- runMega "Image Header" $ do
      dm <- MP.lookAhead MH.parseImageKeywords
      hd <- MH.parseHeader
      pure (dm, hd)
    darr <- mainData dm
    pure $ DataHDU hd darr
  
  
  binTable :: (State ByteString :> es, Parser :> es) => Eff es BinTableHDU
  binTable = do
    (dm, pcount, hd) <- do
      runMega "BinTable Header" $ do
        (dm, pcount) <- lookAhead MH.parseBinTableKeywords
        hd <- MH.parseHeader
        pure (dm, pcount, hd)
  
    darr <- mainData dm
    rest <- get
    let heap = BS.take pcount rest
    put $ BS.dropWhile (== 0) $ BS.drop pcount rest
    pure $ BinTableHDU hd pcount heap darr
  
  
  mainData :: (State ByteString :> es) => Dimensions -> Eff es DataArray
  mainData dm = do
    rest <- get
    let len = dataSizeBytes dm
    let dat = dataArray dm (BS.take len rest)
    put $ BS.dropWhile (== 0) $ BS.drop len rest
    pure dat
  
  
  encodePrimaryHDU :: DataHDU -> ByteString
  encodePrimaryHDU p =
    encodeHDU (renderPrimaryHeader p.header p.dataArray) p.dataArray.rawData
  
  
  encodeImageHDU :: DataHDU -> ByteString
  encodeImageHDU p =
    encodeHDU (renderImageHeader p.header p.dataArray) p.dataArray.rawData
  
  
  encodeExtension :: Extension -> ByteString
  encodeExtension (Image hdu) = encodeImageHDU hdu
  encodeExtension (BinTable _) = error "BinTableHDU rendering not supported"
  
  
  -- | Encode an HDU, properly handling datasum and checksum
  encodeHDU :: (Checksum -> BuilderBlock) -> ByteString -> ByteString
  encodeHDU buildHead rawData =
    let dsum = checksum rawData
     in encodeHeader dsum <> renderDataArray rawData
   where
    encodeHeader :: Checksum -> ByteString
    encodeHeader dsum =
      let h = BS.toStrict $ runRender (buildHead dsum)
          hsum = checksum h -- calculate the checksum of only the header
          csum = hsum <> dsum -- 1s complement add to the datasum
       in replaceChecksum csum h
  
    replaceChecksum :: Checksum -> ByteString -> ByteString
    replaceChecksum csum = replaceKeywordLine "CHECKSUM" (String $ encodeChecksum csum) Nothing
  
  
  -- | Fast replace a single keyword in a raw header bytestring
  replaceKeywordLine :: ByteString -> Value -> Maybe Text -> ByteString -> ByteString
  replaceKeywordLine key val mc header =
    let (start, rest) = BS.breakSubstring key header
        newKeyLine = BS.toStrict $ runRender $ renderKeywordLine (TE.decodeUtf8 key) val mc
     in start <> newKeyLine <> BS.drop 80 rest
  
  
  -- | Parse HDUs by running MegaParsec parsers one at a time, and tracking how much of the ByteString we've consumed
  runMega :: (Parser :> es, State ByteString :> es) => String -> MH.Parser a -> Eff es a
  runMega src parse = do
    inp <- get
    case MH.runNextParser src inp parse of
      Right (a, rest) -> do
        put rest
        pure a
      Left err ->
        parseFail $ MH.showParseError err
  