(*---------------------------------------------------------------------------
   Copyright (c) 2025 Thomas Gazagnaire. All rights reserved.
   SPDX-License-Identifier: ISC
  ---------------------------------------------------------------------------*)

type info = { read_write : bool; sector_size : int; sectors : int64 }
type error = [ `Disconnected | `Read_error of string | `Invalid_sector of int64 ]
type write_error = [ error | `Write_error of string | `Read_only ]

let pp_error ppf = function
  | `Disconnected -> Fmt.string ppf "disconnected"
  | `Read_error s -> Fmt.pf ppf "read error: %s" s
  | `Invalid_sector n -> Fmt.pf ppf "invalid sector: %Ld" n

let pp_write_error ppf = function
  | #error as e -> pp_error ppf e
  | `Write_error s -> Fmt.pf ppf "write error: %s" s
  | `Read_only -> Fmt.string ppf "read-only device"

module type IMPL = sig
  type state

  val info : state -> info
  val read : state -> int64 -> (string, error) result
  val write : state -> int64 -> string -> (unit, write_error) result
  val sync : state -> unit
  val close : state -> unit
end

type t = T : { state : 'a; impl : (module IMPL with type state = 'a) } -> t

let pp ppf (T { state; impl = (module I) }) =
  let i = I.info state in
  Fmt.pf ppf "<block %s sectors=%Ld sector-size=%d>"
    (if i.read_write then "rw" else "ro")
    i.sectors i.sector_size

let info (T { state; impl = (module I) }) = I.info state
let read (T { state; impl = (module I) }) sector = I.read state sector

let write (T { state; impl = (module I) }) sector data =
  I.write state sector data

let sync (T { state; impl = (module I) }) = I.sync state
let close (T { state; impl = (module I) }) = I.close state

let read_many t start count =
  let info = info t in
  let buf = Buffer.create (count * info.sector_size) in
  let rec loop i =
    if i >= count then Ok (Buffer.contents buf)
    else
      match read t Int64.(add start (of_int i)) with
      | Error e -> Error e
      | Ok data ->
          Buffer.add_string buf data;
          loop (i + 1)
  in
  loop 0

let write_many t start sectors =
  let rec loop i = function
    | [] -> Ok ()
    | data :: rest -> (
        match write t Int64.(add start (of_int i)) data with
        | Error e -> Error e
        | Ok () -> loop (i + 1) rest)
  in
  loop 0 sectors

(* Memory implementation *)
module Memory = struct
  type state = {
    mutable data : bytes;
    sector_size : int;
    sectors : int64;
    mutable closed : bool;
  }

  let info t =
    { read_write = true; sector_size = t.sector_size; sectors = t.sectors }

  let read t sector =
    if t.closed then Error `Disconnected
    else if sector < 0L || sector >= t.sectors then
      Error (`Invalid_sector sector)
    else
      let off = Int64.to_int sector * t.sector_size in
      Ok (Bytes.sub_string t.data off t.sector_size)

  let write t sector data =
    if t.closed then Error `Disconnected
    else if sector < 0L || sector >= t.sectors then
      Error (`Invalid_sector sector)
    else if String.length data <> t.sector_size then
      Error (`Write_error "data length must equal sector size")
    else
      let off = Int64.to_int sector * t.sector_size in
      Bytes.blit_string data 0 t.data off t.sector_size;
      Ok ()

  let sync _ = ()
  let close t = t.closed <- true
end

let of_memory ~sector_size ~sectors =
  let size = Int64.to_int sectors * sector_size in
  let state =
    {
      Memory.data = Bytes.make size '\x00';
      sector_size;
      sectors;
      closed = false;
    }
  in
  T { state; impl = (module Memory) }

let of_bigarray ~sector_size ba =
  let len = Bigarray.Array1.dim ba in
  if len mod sector_size <> 0 then
    invalid_arg "bigarray size must be multiple of sector_size";
  let sectors = Int64.of_int (len / sector_size) in
  let data = Bytes.create len in
  for i = 0 to len - 1 do
    Bytes.set data i (Bigarray.Array1.get ba i)
  done;
  let state = { Memory.data; sector_size; sectors; closed = false } in
  T { state; impl = (module Memory) }

(* Read-only wrapper *)
module Read_only = struct
  type state = { inner : t }

  let info t =
    let i = info t.inner in
    { i with read_write = false }

  let read t sector = read t.inner sector
  let write _ _ _ = Error `Read_only
  let sync t = sync t.inner
  let close t = close t.inner
end

let read_only t =
  let state = { Read_only.inner = t } in
  T { state; impl = (module Read_only) }

let of_string ~sector_size data =
  let len = String.length data in
  if len mod sector_size <> 0 then
    invalid_arg "string length must be multiple of sector_size";
  let sectors = Int64.of_int (len / sector_size) in
  let state =
    { Memory.data = Bytes.of_string data; sector_size; sectors; closed = false }
  in
  (* Return read-only by wrapping *)
  let t = T { state; impl = (module Memory) } in
  read_only t

(* File implementation *)
module File = struct
  type state = {
    file : Eio.File.rw_ty Eio.Resource.t;
    sector_size : int;
    sectors : int64;
    mutable closed : bool;
  }

  let info t =
    { read_write = true; sector_size = t.sector_size; sectors = t.sectors }

  let read t sector =
    if t.closed then Error `Disconnected
    else if sector < 0L || sector >= t.sectors then
      Error (`Invalid_sector sector)
    else
      try
        let off =
          Optint.Int63.of_int64 (Int64.mul sector (Int64.of_int t.sector_size))
        in
        let buf = Cstruct.create t.sector_size in
        Eio.File.pread_exact t.file ~file_offset:off [ buf ];
        Ok (Cstruct.to_string buf)
      with exn -> Error (`Read_error (Printexc.to_string exn))

  let write t sector data =
    if t.closed then Error `Disconnected
    else if sector < 0L || sector >= t.sectors then
      Error (`Invalid_sector sector)
    else if String.length data <> t.sector_size then
      Error (`Write_error "data length must equal sector size")
    else
      try
        let off =
          Optint.Int63.of_int64 (Int64.mul sector (Int64.of_int t.sector_size))
        in
        let buf = Cstruct.of_string data in
        Eio.File.pwrite_all t.file ~file_offset:off [ buf ];
        Ok ()
      with exn -> Error (`Write_error (Printexc.to_string exn))

  let sync t = if not t.closed then Eio.File.sync t.file

  let close t =
    if not t.closed then (
      t.closed <- true;
      Eio.Resource.close t.file)
end

let of_file ~sw path ~sector_size ?sectors:sectors_override ?create () =
  if sector_size < 512 || sector_size land (sector_size - 1) <> 0 then
    invalid_arg "sector_size must be power of 2 and >= 512";
  let file =
    match create with
    | Some sectors ->
        let size = Int64.mul sectors (Int64.of_int sector_size) in
        let f = Eio.Path.open_out ~sw path ~create:(`Or_truncate 0o644) in
        (* Extend file to size *)
        Eio.File.truncate f (Optint.Int63.of_int64 size);
        (f :> Eio.File.rw_ty Eio.Resource.t)
    | None ->
        Eio.Path.open_out ~sw path ~create:`Never |> fun f ->
        (f :> Eio.File.rw_ty Eio.Resource.t)
  in
  let sectors =
    match sectors_override with
    | Some s -> s
    | None ->
        let stat = Eio.File.stat file in
        let size = Optint.Int63.to_int64 stat.size in
        Int64.div size (Int64.of_int sector_size)
  in
  let state = { File.file; sector_size; sectors; closed = false } in
  T { state; impl = (module File) }

(* Sub-device wrapper *)
module Sub = struct
  type state = { inner : t; start : int64; sectors : int64 }

  let info t =
    let i = info t.inner in
    { i with sectors = t.sectors }

  let read t sector =
    if sector < 0L || sector >= t.sectors then Error (`Invalid_sector sector)
    else read t.inner Int64.(add t.start sector)

  let write t sector data =
    if sector < 0L || sector >= t.sectors then Error (`Invalid_sector sector)
    else write t.inner Int64.(add t.start sector) data

  let sync t = sync t.inner
  let close _ = () (* Don't close underlying device *)
end

let sub t ~start ~sectors =
  let i = info t in
  if start < 0L || sectors < 0L || Int64.add start sectors > i.sectors then
    invalid_arg "sub: invalid range";
  let state = { Sub.inner = t; start; sectors } in
  T { state; impl = (module Sub) }

(* CRC32C wrapper *)
module With_crc = struct
  type state = { inner : t; data_size : int }

  let crc32c data =
    (* Simple CRC32C implementation - in production use Wal.crc32c *)
    let crc = ref 0xFFFFFFFFl in
    for i = 0 to String.length data - 1 do
      let byte = Char.code data.[i] in
      crc := Int32.logxor !crc (Int32.of_int byte);
      for _ = 0 to 7 do
        let mask = Int32.neg (Int32.logand !crc 1l) in
        crc :=
          Int32.logxor
            (Int32.shift_right_logical !crc 1)
            (Int32.logand 0x82F63B78l mask)
      done
    done;
    Int32.logxor !crc 0xFFFFFFFFl

  let encode_crc crc =
    let b = Bytes.create 4 in
    Bytes.set b 0 (Char.chr (Int32.to_int (Int32.logand crc 0xFFl)));
    Bytes.set b 1
      (Char.chr
         (Int32.to_int (Int32.logand (Int32.shift_right_logical crc 8) 0xFFl)));
    Bytes.set b 2
      (Char.chr
         (Int32.to_int (Int32.logand (Int32.shift_right_logical crc 16) 0xFFl)));
    Bytes.set b 3 (Char.chr (Int32.to_int (Int32.shift_right_logical crc 24)));
    Bytes.to_string b

  let decode_crc s off =
    let b0 = Int32.of_int (Char.code s.[off]) in
    let b1 = Int32.of_int (Char.code s.[off + 1]) in
    let b2 = Int32.of_int (Char.code s.[off + 2]) in
    let b3 = Int32.of_int (Char.code s.[off + 3]) in
    Int32.logor b0
      (Int32.logor (Int32.shift_left b1 8)
         (Int32.logor (Int32.shift_left b2 16) (Int32.shift_left b3 24)))

  let info t =
    let i = info t.inner in
    { i with sector_size = t.data_size }

  let read t sector =
    match read t.inner sector with
    | Error e -> Error e
    | Ok raw ->
        let data = String.sub raw 0 t.data_size in
        let stored_crc = decode_crc raw t.data_size in
        let computed_crc = crc32c data in
        if stored_crc <> computed_crc then
          Error
            (`Read_error
               (Fmt.str "CRC mismatch: stored=%lx computed=%lx" stored_crc
                  computed_crc))
        else Ok data

  let write t sector data =
    if String.length data <> t.data_size then
      Error (`Write_error (Fmt.str "data length must be %d" t.data_size))
    else
      let crc = crc32c data in
      let raw = data ^ encode_crc crc in
      write t.inner sector raw

  let sync t = sync t.inner
  let close t = close t.inner
end

let with_crc32c t =
  let i = info t in
  let data_size = i.sector_size - 4 in
  if data_size < 1 then invalid_arg "sector_size too small for CRC";
  let state = { With_crc.inner = t; data_size } in
  T { state; impl = (module With_crc) }

(* Bytesrw integration *)
let to_reader t ~offset ~length =
  let i = info t in
  if Int64.rem offset (Int64.of_int i.sector_size) <> 0L then
    invalid_arg "offset must be sector-aligned";
  let start_sector = Int64.div offset (Int64.of_int i.sector_size) in
  let current_sector = ref start_sector in
  let current_pos = ref 0 in
  let remaining = ref length in
  let current_data = ref "" in
  Bytesrw.Bytes.Reader.make ~slice_length:i.sector_size @@ fun () ->
  if !remaining <= 0L then Bytesrw.Bytes.Slice.eod
  else if !current_pos >= String.length !current_data then (
    (* Read next sector *)
    match read t !current_sector with
    | Error _ -> Bytesrw.Bytes.Slice.eod
    | Ok data ->
        current_data := data;
        current_pos := 0;
        current_sector := Int64.succ !current_sector;
        let len = min (String.length data) (Int64.to_int !remaining) in
        remaining := Int64.sub !remaining (Int64.of_int len);
        current_pos := len;
        Bytesrw.Bytes.Slice.make
          (Bytes.unsafe_of_string data)
          ~first:0 ~length:len)
  else Bytesrw.Bytes.Slice.eod

let to_writer t ~offset =
  let i = info t in
  if Int64.rem offset (Int64.of_int i.sector_size) <> 0L then
    invalid_arg "offset must be sector-aligned";
  let start_sector = Int64.div offset (Int64.of_int i.sector_size) in
  let current_sector = ref start_sector in
  let buffer = Buffer.create i.sector_size in
  let flush () =
    if Buffer.length buffer > 0 then (
      (* Pad to sector size if needed *)
      while Buffer.length buffer < i.sector_size do
        Buffer.add_char buffer '\x00'
      done;
      let _ = write t !current_sector (Buffer.contents buffer) in
      Buffer.clear buffer;
      current_sector := Int64.succ !current_sector)
  in
  Bytesrw.Bytes.Writer.make ~slice_length:i.sector_size @@ fun slice ->
  if Bytesrw.Bytes.Slice.is_eod slice then flush ()
  else begin
    let data =
      Bytes.sub_string
        (Bytesrw.Bytes.Slice.bytes slice)
        (Bytesrw.Bytes.Slice.first slice)
        (Bytesrw.Bytes.Slice.length slice)
    in
    Buffer.add_string buffer data;
    while Buffer.length buffer >= i.sector_size do
      let sector_data = Buffer.sub buffer 0 i.sector_size in
      let _ = write t !current_sector sector_data in
      current_sector := Int64.succ !current_sector;
      let remaining =
        Buffer.sub buffer i.sector_size (Buffer.length buffer - i.sector_size)
      in
      Buffer.clear buffer;
      Buffer.add_string buffer remaining
    done
  end

(* Generic operations *)

let fold ~f t init =
  let i = info t in
  let rec loop sector acc =
    if sector >= i.sectors then Ok acc
    else
      match read t sector with
      | Error e -> Error e
      | Ok data -> (
          match f sector data acc with
          | Error e -> Error e
          | Ok acc' -> loop (Int64.succ sector) acc')
  in
  loop 0L init

let iter ~f t = fold ~f:(fun sector data () -> f sector data) t ()

type compare_error =
  [ error
  | `Different_sizes
  | `Different_sector_sizes
  | `Contents_differ of int64 ]

let pp_compare_error ppf = function
  | #error as e -> pp_error ppf e
  | `Different_sizes -> Fmt.string ppf "different sizes"
  | `Different_sector_sizes -> Fmt.string ppf "different sector sizes"
  | `Contents_differ n -> Fmt.pf ppf "contents differ at sector %Ld" n

let compare a b =
  let ia = info a and ib = info b in
  if ia.sector_size <> ib.sector_size then Error `Different_sector_sizes
  else if ia.sectors <> ib.sectors then Error `Different_sizes
  else
    let rec loop sector =
      if sector >= ia.sectors then Ok ()
      else
        match (read a sector, read b sector) with
        | Error e, _ -> Error (e :> compare_error)
        | _, Error e -> Error (e :> compare_error)
        | Ok da, Ok db ->
            if da <> db then Error (`Contents_differ sector)
            else loop (Int64.succ sector)
    in
    loop 0L

type copy_error = [ write_error | `Different_sizes | `Different_sector_sizes ]

let pp_copy_error ppf = function
  | #write_error as e -> pp_write_error ppf e
  | `Different_sizes -> Fmt.string ppf "different sizes"
  | `Different_sector_sizes -> Fmt.string ppf "different sector sizes"

let copy ~src ~dst =
  let is = info src and id = info dst in
  if is.sector_size <> id.sector_size then Error `Different_sector_sizes
  else if is.sectors > id.sectors then Error `Different_sizes
  else
    let rec loop sector =
      if sector >= is.sectors then Ok ()
      else
        match read src sector with
        | Error e -> Error (e :> copy_error)
        | Ok data -> (
            match write dst sector data with
            | Error e -> Error (e :> copy_error)
            | Ok () -> loop (Int64.succ sector))
    in
    loop 0L

let is_zero data =
  let rec loop i =
    if i >= String.length data then true
    else if data.[i] <> '\x00' then false
    else loop (i + 1)
  in
  loop 0

let sparse_copy ~src ~dst =
  let is = info src and id = info dst in
  if is.sector_size <> id.sector_size then Error `Different_sector_sizes
  else if is.sectors > id.sectors then Error `Different_sizes
  else
    let rec loop sector =
      if sector >= is.sectors then Ok ()
      else
        match read src sector with
        | Error e -> Error (e :> copy_error)
        | Ok data -> (
            if is_zero data then loop (Int64.succ sector)
            else
              match write dst sector data with
              | Error e -> Error (e :> copy_error)
              | Ok () -> loop (Int64.succ sector))
    in
    loop 0L

let fill t c =
  let i = info t in
  let data = String.make i.sector_size c in
  let rec loop sector =
    if sector >= i.sectors then Ok ()
    else
      match write t sector data with
      | Error e -> Error e
      | Ok () -> loop (Int64.succ sector)
  in
  loop 0L

let zero t = fill t '\x00'