import gleam/bit_array
import gleam/erlang/process.{type Subject}
import gleam/list
import gleam/otp/actor
import gleam/otp/supervision

/// Compress some data into Zstandard format.
///
/// Returns an error if the data isn't byte-aligned.
@external(erlang, "regular_ffi", "compress")
pub fn compress(data: BitArray) -> Result(BitArray, Nil)

/// Decompress some Zstandard data.
///
/// Returns an error if the data isn't valid Zstd data or if it isn't byte-aligned.
@external(erlang, "regular_ffi", "decompress")
pub fn decompress(data: BitArray) -> Result(BitArray, Nil)

type CompressionContext

@external(erlang, "regular_ffi", "new_compression_context")
fn new_compression_context() -> CompressionContext

type Compressor {
  Compressor(context: CompressionContext)
}

pub opaque type CompressorMessage {
  CompressorStream(
    chunk: BitArray,
    respond_to: Subject(Result(CompressionResult, Nil)),
  )
  CompressorFinishFrame(
    chunk: BitArray,
    respond_to: Subject(Result(BitArray, Nil)),
  )
  CompressorShutdown
}

/// Starts a new compressor, used to create streamed compressed data.
///
/// It is encouraged to use [`supervised_compressor`](#supervised_compressor) over this function.
pub fn start_compressor() -> Result(
  actor.Started(Subject(CompressorMessage)),
  actor.StartError,
) {
  actor.new_with_initialiser(5000, fn(self) {
    let context = new_compression_context()

    let selector =
      process.new_selector()
      |> process.select(self)

    actor.initialised(Compressor(context))
    |> actor.selecting(selector)
    |> actor.returning(self)
    |> Ok
  })
  |> actor.on_message(on_compressor_message)
  |> actor.start
}

/// Starts a new supervised compressor, used to create streamed compressed data.
pub fn supervised_compressor() -> supervision.ChildSpecification(
  Subject(CompressorMessage),
) {
  supervision.worker(start_compressor)
}

pub type CompressionResult {
  /// The provided chunk has been fully compressed.
  ///
  /// Note: The output MAY be empty and returned upon calling `finish_compressor_frame`.
  CompressionFinished(output: BitArray)
  /// The provided chunk has been partially compressed.
  ///
  /// Run `stream_compressor` on the remainder to try to compress the full chunk.
  CompressionUnfinished(output: BitArray, remainder: BitArray)
}

/// Stream data to the compressor. The compressor will only compress ~128kB of data
/// at a time.
///
/// If the compressor fails to compress all of the data in one go, `Ok(CompressionUnfinished(..))` will be returned.
/// Call `stream_to_compressor` on the remainder to attempt compressing the rest of the chunk.
///
/// Do so before continuing to the next chunk, otherwise you will end up with malformed data.
/// ```
/// chunk1 -> chunk1's remainders -> chunk2
/// ```
///
/// Returns `Error(Nil)` if the data isn't byte-aligned.
pub fn stream_to_compressor(
  compressor: Subject(CompressorMessage),
  waiting timeout: Int,
  chunk chunk: BitArray,
) -> Result(CompressionResult, Nil) {
  process.call(compressor, waiting: timeout, sending: CompressorStream(
    chunk:,
    respond_to: _,
  ))
}

/// Stream data to the compressor, continuing to compress until the end of the chunk.
///
/// Timeout applies to each subsequent call to `stream_to_compressor` - this function just calls
/// that one until it finishes the chunk.
/// 
/// Enforce size limits on untrusted data!
pub fn stream_all_to_compressor(
  compressor: Subject(CompressorMessage),
  waiting timeout: Int,
  chunk chunk: BitArray,
) -> Result(BitArray, Nil) {
  stream_all_to_compressor_loop(compressor, timeout, chunk, [])
}

fn stream_all_to_compressor_loop(
  compressor: Subject(CompressorMessage),
  timeout: Int,
  chunk: BitArray,
  acc: List(BitArray),
) -> Result(BitArray, Nil) {
  let result = stream_to_compressor(compressor, waiting: timeout, chunk:)

  case result {
    Ok(CompressionFinished(output)) ->
      [output, ..acc]
      |> list.reverse
      |> bit_array.concat
      |> Ok
    Ok(CompressionUnfinished(output, remainder)) ->
      stream_all_to_compressor_loop(compressor, timeout, remainder, [
        output,
        ..acc
      ])
    Error(Nil) -> Error(Nil)
  }
}

pub fn shutdown_compressor(compressor: Subject(CompressorMessage)) -> Nil {
  process.send(compressor, CompressorShutdown)
}

/// Flushes the compressor's output buffer, making it ready to compress new data.
///
/// Returns the final compressed chunk, if successful.
///
/// Returns `Error(Nil)` if the data isn't byte-aligned.
pub fn finish_compressor_frame(
  compressor: Subject(CompressorMessage),
  waiting timeout: Int,
  final_chunk chunk: BitArray,
) -> Result(BitArray, Nil) {
  process.call(compressor, waiting: timeout, sending: CompressorFinishFrame(
    chunk:,
    respond_to: _,
  ))
}

fn on_compressor_message(compressor: Compressor, message: CompressorMessage) {
  case message {
    CompressorStream(chunk:, respond_to:) -> {
      let result = compress_chunk(compressor.context, chunk)
      process.send(respond_to, result)
      actor.continue(compressor)
    }
    CompressorFinishFrame(chunk:, respond_to:) -> {
      let result = finish_compress_frame(compressor.context, chunk)
      process.send(respond_to, result)
      actor.continue(compressor)
    }
    CompressorShutdown -> actor.stop()
  }
}

@external(erlang, "regular_ffi", "compress_chunk")
fn compress_chunk(
  context: CompressionContext,
  chunk: BitArray,
) -> Result(CompressionResult, Nil)

@external(erlang, "regular_ffi", "finish_compress_frame")
fn finish_compress_frame(
  context: CompressionContext,
  chunk: BitArray,
) -> Result(BitArray, Nil)

type DecompressionContext

@external(erlang, "regular_ffi", "new_decompression_context")
fn new_decompression_context() -> DecompressionContext

type Decompressor {
  Decompressor(context: DecompressionContext)
}

pub opaque type DecompressorMessage {
  DecompressorStream(
    chunk: BitArray,
    respond_to: Subject(Result(DecompressionResult, Nil)),
  )
  DecompressorFinishFrame(
    chunk: BitArray,
    respond_to: Subject(Result(BitArray, Nil)),
  )
  DecompressorShutdown
}

/// Starts a new decompressor, used to decompress data flowing in through a stream.
///
/// It is encouraged to use [`supervised_decompressor`](#supervised_decompressor) over this function.
pub fn start_decompressor() -> Result(
  actor.Started(Subject(DecompressorMessage)),
  actor.StartError,
) {
  actor.new_with_initialiser(5000, fn(self) {
    let context = new_decompression_context()

    let selector =
      process.new_selector()
      |> process.select(self)

    actor.initialised(Decompressor(context))
    |> actor.selecting(selector)
    |> actor.returning(self)
    |> Ok
  })
  |> actor.on_message(on_decompressor_message)
  |> actor.start
}

/// Starts a new supervised decompressor, used to decompress data flowing in through a stream.
pub fn supervised_decompressor() -> supervision.ChildSpecification(
  Subject(DecompressorMessage),
) {
  supervision.worker(start_decompressor)
}

pub type DecompressionResult {
  /// The provided chunk has been fully decompressed.
  ///
  /// Note: The output MAY be empty and returned upon calling `finish_decompressor_frame`.
  DecompressionFinished(output: BitArray)
  /// The provided chunk has been partially decompressed.
  ///
  /// Run `stream_decompressor` on the remainder to try to decompress the full chunk.
  DecompressionUnfinished(output: BitArray, remainder: BitArray)
}

/// Stream data to the decompressor. The compressor will only inflate ~128kB of data
/// at a time.
///
/// If the decompressor fails to inflate all of the data in one go, `Ok(DecompressionUnfinished(..))` will be returned.
/// Call `stream_to_decompressor` on the remainder to attempt compressing the rest of the chunk.
///
/// Do so before continuing to the next chunk, otherwise you will end up with malformed inflated data.
/// ```
/// chunk1 -> chunk1's remainders -> chunk2
/// ```
///
/// Returns `Error(Nil)` if the data isn't byte-aligned or isn't valid zstd data.
pub fn stream_to_decompressor(
  compressor: Subject(DecompressorMessage),
  waiting timeout: Int,
  chunk chunk: BitArray,
) -> Result(DecompressionResult, Nil) {
  process.call(compressor, waiting: timeout, sending: DecompressorStream(
    chunk:,
    respond_to: _,
  ))
}

/// Stream data to the decompressor, continuing to inflate until the end of the chunk.
///
/// Timeout applies to each subsequent call to `stream_to_decompressor` - this function just calls
/// that one until it finishes the chunk.
/// 
/// Enforce size limits on untrusted data!
pub fn stream_all_to_decompressor(
  compressor: Subject(DecompressorMessage),
  waiting timeout: Int,
  chunk chunk: BitArray,
) -> Result(BitArray, Nil) {
  stream_all_to_decompressor_loop(compressor, timeout, chunk, [])
}

fn stream_all_to_decompressor_loop(
  compressor: Subject(DecompressorMessage),
  timeout: Int,
  chunk: BitArray,
  acc: List(BitArray),
) -> Result(BitArray, Nil) {
  let result = stream_to_decompressor(compressor, waiting: timeout, chunk:)

  case result {
    Ok(DecompressionFinished(output)) ->
      [output, ..acc]
      |> list.reverse
      |> bit_array.concat
      |> Ok
    Ok(DecompressionUnfinished(output, remainder)) ->
      stream_all_to_decompressor_loop(compressor, timeout, remainder, [
        output,
        ..acc
      ])
    Error(Nil) -> Error(Nil)
  }
}

pub fn shutdown_decompressor(compressor: Subject(DecompressorMessage)) -> Nil {
  process.send(compressor, DecompressorShutdown)
}

/// Flushes the decompressor's output buffer, making it ready to decompress new data.
///
/// Returns the final inflated chunk, if successful.
///
/// Returns `Error(Nil)` if the data isn't byte-aligned or isn't valid zstd data.
pub fn finish_decompressor_frame(
  compressor: Subject(DecompressorMessage),
  waiting timeout: Int,
  final_chunk chunk: BitArray,
) -> Result(BitArray, Nil) {
  process.call(compressor, waiting: timeout, sending: DecompressorFinishFrame(
    chunk:,
    respond_to: _,
  ))
}

fn on_decompressor_message(
  compressor: Decompressor,
  message: DecompressorMessage,
) {
  case message {
    DecompressorStream(chunk:, respond_to:) -> {
      let result = decompress_chunk(compressor.context, chunk)
      process.send(respond_to, result)
      actor.continue(compressor)
    }
    DecompressorFinishFrame(chunk:, respond_to:) -> {
      let result = finish_decompress_frame(compressor.context, chunk)
      process.send(respond_to, result)
      actor.continue(compressor)
    }
    DecompressorShutdown -> actor.stop()
  }
}

@external(erlang, "regular_ffi", "decompress_chunk")
fn decompress_chunk(
  context: DecompressionContext,
  chunk: BitArray,
) -> Result(DecompressionResult, Nil)

@external(erlang, "regular_ffi", "finish_decompress_frame")
fn finish_decompress_frame(
  context: DecompressionContext,
  chunk: BitArray,
) -> Result(BitArray, Nil)