module Wire = struct
  let split_once ch value =
    match String.index_opt value ch with
    | None -> None
    | Some index ->
        Some
          ( String.sub value 0 index,
            String.sub value (index + 1) (String.length value - index - 1) )

  let lines bytes =
    bytes |> Bytes.to_string |> String.split_on_char '\n'
    |> List.filter (fun line -> String.trim line <> "")

  let fields bytes =
    lines bytes
    |> List.filter_map (fun line -> split_once '=' line)
    |> List.map (fun (key, value) -> (String.trim key, String.trim value))

  let require_type expected fields =
    match List.assoc_opt "type" fields with
    | Some actual when actual = expected -> ()
    | Some actual ->
        invalid_arg
          (Printf.sprintf "expected storable type %S, got %S" expected actual)
    | None -> invalid_arg "missing storable type"

  let values key fields =
    List.filter_map (fun (k, v) -> if k = key then Some v else None) fields

  let required key fields =
    match List.assoc_opt key fields with
    | Some value -> value
    | None -> invalid_arg ("missing storable field: " ^ key)

  let attribute value =
    match split_once ':' value with
    | Some pair -> pair
    | None -> invalid_arg ("invalid attribute field: " ^ value)

  let encode fields =
    fields
    |> List.map (fun (key, value) -> key ^ "=" ^ value)
    |> String.concat "\n" |> fun value -> Bytes.of_string (value ^ "\n")
end

module Tuple = struct
  type t = { relation : string; attributes : (string * string) list }

  let to_bytes tuple =
    Wire.encode
      ([ ("type", "tuple"); ("relation", tuple.relation) ]
       @ List.map
           (fun (name, value) -> ("attribute", name ^ ":" ^ value))
           tuple.attributes)

  let of_bytes bytes =
    let fields = Wire.fields bytes in
    Wire.require_type "tuple" fields;
    {
      relation = Wire.required "relation" fields;
      attributes = List.map Wire.attribute (Wire.values "attribute" fields);
    }
end

module Relation = struct
  type t = {
    name : string;
    schema : (string * string) list;
    (** Hex Merkle root of the relation's tuple set.  Empty string = empty relation. *)
    tree_pointer : string;
  }

  let to_bytes relation =
    Wire.encode
      ([ ("type", "relation"); ("name", relation.name); ("kind", "stored");
         ("tree_pointer", relation.tree_pointer) ]
       @ List.map
           (fun (name, domain) -> ("attribute", name ^ ":" ^ domain))
           relation.schema)

  let of_bytes bytes =
    let fields = Wire.fields bytes in
    Wire.require_type "relation" fields;
    {
      name         = Wire.required "name" fields;
      schema       = List.map Wire.attribute (Wire.values "attribute" fields);
      tree_pointer = Option.value (List.assoc_opt "tree_pointer" fields) ~default:"";
    }
end

module Multigroup = struct
  type t = { name : string; relations : Relation.t list }

  (* Wire format: "name|attr1:dom1,attr2:dom2|<merkle_root>"
     The third segment is the Merkle root hex; empty string = empty relation.
     A fourth segment for constraints is reserved but not yet implemented.
     TODO: add constraints serialisation once a canonical sexp representation
     is agreed upon (two syntactically distinct bodies may be semantically
     equivalent — the round-trip must account for that). *)
  let relation_value relation =
    let attrs =
      relation.Relation.schema
      |> List.map (fun (name, domain) -> name ^ ":" ^ domain)
      |> String.concat ","
    in
    relation.name ^ "|" ^ attrs ^ "|" ^ relation.tree_pointer

  let parse_relation value =
    match Wire.split_once '|' value with
    | None -> { Relation.name = value; schema = []; tree_pointer = "" }
    | Some (name, rest) ->
        let attrs_str, tree_pointer =
          match Wire.split_once '|' rest with
          | None         -> (rest, "")
          | Some (a, tp) -> (a, tp)
        in
        let schema =
          if attrs_str = "" then []
          else attrs_str |> String.split_on_char ',' |> List.map Wire.attribute
        in
        { Relation.name; schema; tree_pointer }

  let to_bytes multigroup =
    Wire.encode
      ([ ("type", "multigroup"); ("name", multigroup.name) ]
       @ List.map
           (fun relation -> ("relation", relation_value relation))
           multigroup.relations)

  let of_bytes bytes =
    let fields = Wire.fields bytes in
    Wire.require_type "multigroup" fields;
    {
      name = Wire.required "name" fields;
      relations = List.map parse_relation (Wire.values "relation" fields);
    }
end