(**************************************************************************)
(*                                BELENIOS                                *)
(*                                                                        *)
(*  Copyright © 2012-2022 Inria                                           *)
(*                                                                        *)
(*  This program is free software: you can redistribute it and/or modify  *)
(*  it under the terms of the GNU Affero General Public License as        *)
(*  published by the Free Software Foundation, either version 3 of the    *)
(*  License, or (at your option) any later version, with the additional   *)
(*  exemption that compiling, linking, and/or using OpenSSL is allowed.   *)
(*                                                                        *)
(*  This program is distributed in the hope that it will be useful, but   *)
(*  WITHOUT ANY WARRANTY; without even the implied warranty of            *)
(*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU     *)
(*  Affero General Public License for more details.                       *)
(*                                                                        *)
(*  You should have received a copy of the GNU Affero General Public      *)
(*  License along with this program.  If not, see                         *)
(*  <http://www.gnu.org/licenses/>.                                       *)
(**************************************************************************)

open Belenios_core.Serializable_j
open Belenios_core.Common
open Belenios_core.Events
module Archive = Belenios_core.Archive

let block_size = Archive.block_size

type index = {
  map : (hash, location) Hashtbl.t;
  mutable roots : roots;
  mutable last_event : event option;
  file : string;
  mutable lines : (Archive.data_or_event * hash) list;
  timestamp : int64;
  header : archive_header;
}

module DirectMonad = struct
  type 'a t = 'a

  let return x = x
  let bind x f = f x
  let fail x = raise x
  let yield () = ()
end

module IoReader = struct
  include DirectMonad

  type file = in_channel

  let get_pos = LargeFile.pos_in
  let set_pos = LargeFile.seek_in
  let read_block ic buffer = really_input ic buffer 0 block_size
end

module Reader = Archive.MakeReader (IoReader)

module IoWriter = struct
  include DirectMonad

  type file = out_channel

  let get_pos = LargeFile.pos_out
  let write_block oc buffer = output_bytes oc buffer
end

module Writer = Archive.MakeWriter (IoWriter)

let build_index filename =
  let r = Hashtbl.create 1000 in
  let ic = open_in filename in
  let header = Reader.read_header ic in
  let rec loop last accu lines =
    match Reader.read_record ic with
    | exception End_of_file -> (r, last, accu, lines, header)
    | record ->
        let last, accu =
          match record.typ with
          | Data -> (last, accu)
          | Event event -> (Some event, update_roots record.hash event accu)
        in
        Hashtbl.add r record.hash record.location;
        loop last accu ((record.typ, record.hash) :: lines)
  in
  Fun.protect
    ~finally:(fun () -> close_in ic)
    (fun () -> loop None empty_roots [])

let get_index ~file =
  let map, last_event, roots, lines, header = build_index file in
  let timestamp = Archive.get_timestamp header in
  { map; roots; last_event; file; lines; timestamp; header }

let gethash ~index ~filename x =
  match Hashtbl.find_opt index x with
  | None -> None
  | Some i ->
      let ic = open_in filename in
      Fun.protect
        ~finally:(fun () -> close_in ic)
        (fun () ->
          LargeFile.seek_in ic i.location_offset;
          assert (i.location_length <= Int64.of_int Sys.max_string_length);
          Some (really_input_string ic (Int64.to_int i.location_length)))

let get_data i x = gethash ~index:i.map ~filename:i.file x

let get_event i x =
  gethash ~index:i.map ~filename:i.file x |> Option.map event_of_string

let get_roots i = i.roots

let fold_on_event_payload_hashes index typ last_event f accu =
  let rec loop e accu =
    match get_event index e with
    | None -> assert false
    | Some e ->
        if e.event_typ = typ then
          match (e.event_payload, e.event_parent) with
          | Some payload, Some parent -> loop parent (f payload accu)
          | _ -> assert false
        else accu
  in
  loop last_event accu

let fold_on_event_payloads index typ last_event f accu =
  fold_on_event_payload_hashes index typ last_event
    (fun payload accu ->
      match get_data index payload with
      | None -> assert false
      | Some x -> f x accu)
    accu

let fsck index =
  let last_event =
    match index.last_event with None -> failwith "no events" | Some x -> x
  in
  let module IoComparer = struct
    include DirectMonad

    type file = in_channel

    let get_pos = LargeFile.pos_in
    let buffer = Bytes.create block_size

    let write_block ic expected =
      try
        let () = really_input ic buffer 0 block_size in
        if expected <> buffer then
          failwith "generated archive is not identical to original one"
      with End_of_file ->
        failwith "generate archive is longer than original one"
  end in
  let module Comparer = Archive.MakeWriter (IoComparer) in
  let module IoArchiver = struct
    include DirectMonad

    let get_hash hash = gethash ~index:index.map ~filename:index.file hash
  end in
  let module Archiver = Archive.MakeArchiver (IoArchiver) (Comparer) in
  let ic = open_in_bin index.file in
  let length = LargeFile.in_channel_length ic in
  Fun.protect
    ~finally:(fun () -> close_in ic)
    (fun () ->
      let () = Archiver.write_archive ic index.header last_event in
      if LargeFile.pos_in ic <> length then
        failwith "generated archive is shorter than original one")

let starts_with ~(prefix : index) (index : index) =
  let rec loop x y =
    match (x, y) with
    | x :: xs, y :: ys when x = y -> loop xs ys
    | [], _ -> true
    | _ -> false
  in
  loop (List.rev prefix.lines) (List.rev index.lines)

let write_header filename header =
  let oc =
    open_out_gen
      [ Open_wronly; Open_append; Open_creat; Open_binary ]
      0o644 filename
  in
  Fun.protect
    ~finally:(fun () -> close_out oc)
    (fun () -> Writer.write_header oc header)

let raw_append ~filename ~timestamp xs =
  let oc =
    open_out_gen [ Open_wronly; Open_append; Open_binary ] 0o644 filename
  in
  Fun.protect
    ~finally:(fun () -> close_out oc)
    (fun () ->
      List.fold_left
        (fun accu (typ, x) -> Writer.write_record oc ~timestamp typ x :: accu)
        [] xs)

type append_operation = Data of string | Event of event_type * hash option

let append index ops =
  let last_event, roots, items, lines =
    List.fold_left
      (fun (last_event, roots, items, lines) x ->
        match x with
        | Data x ->
            let typ = Archive.Data in
            let items = (typ, x) :: items in
            let lines = (typ, Hash.hash_string x) :: lines in
            (last_event, roots, items, lines)
        | Event (event_typ, event_payload) ->
            let event_parent, event_height =
              match last_event with
              | None -> (None, 0)
              | Some x ->
                  ( Some (Hash.hash_string (string_of_event x)),
                    x.event_height + 1 )
            in
            let event =
              { event_parent; event_height; event_typ; event_payload }
            in
            let typ = Archive.Event event in
            let event_s = string_of_event event in
            let event_h = Hash.hash_string event_s in
            let roots = update_roots event_h event roots in
            let items = (typ, event_s) :: items in
            (Some event, roots, items, (typ, event_h) :: lines))
      (index.last_event, index.roots, [], index.lines)
      ops
  in
  let items = List.rev items in
  let records =
    raw_append ~filename:index.file ~timestamp:index.timestamp items
  in
  List.iter (fun r -> Hashtbl.add index.map r.Archive.hash r.location) records;
  index.roots <- roots;
  index.last_event <- last_event;
  index.lines <- lines

let init ~file ~election ~trustees ~public_creds =
  if Sys.file_exists file then Printf.ksprintf failwith "%s already exists" file;
  let header = Archive.new_header () in
  let index =
    {
      map = Hashtbl.create 1000;
      roots = empty_roots;
      last_event = None;
      file;
      lines = [];
      timestamp = Archive.get_timestamp header;
      header;
    }
  in
  write_header file header;
  let setup_election = Hash.hash_string election in
  let setup_trustees = Hash.hash_string trustees in
  let setup_credentials = Hash.hash_string public_creds in
  let setup_data = { setup_election; setup_trustees; setup_credentials } in
  let setup_data_s = string_of_setup_data setup_data in
  append index
    [
      Data election;
      Data trustees;
      Data public_creds;
      Data setup_data_s;
      Event (`Setup, Some (Hash.hash_string setup_data_s));
    ];
  index
