open Hito_app
module WI = Workout_import

type error = Malformed of string | Unknown_exercise of string | Invalid_time

let pp_error ppf = function
  | Malformed detail -> Format.fprintf ppf "malformed import record: %s" detail
  | Unknown_exercise id -> Format.fprintf ppf "unknown exercise %S" id
  | Invalid_time -> Format.fprintf ppf "ended_at precedes started_at"

(* --- length-prefixed writer ---

   The writer emits every field as its byte length in decimal, a colon, then
   its raw bytes. The length frames the bytes exactly. A field may therefore
   hold any byte, including newlines, commas, and colons, and the reader never
   has to escape or guess where a field ends. *)

module Writer = struct
  let create () = Buffer.create 256

  let field buf (s : string) =
    Buffer.add_string buf (string_of_int (String.length s));
    Buffer.add_char buf ':';
    Buffer.add_string buf s

  let string = field
  let int buf n = field buf (string_of_int n)

  (* [%h] renders a float as a round-trippable hexadecimal literal. *)
  let float buf x = field buf (Printf.sprintf "%h" x)
  let timestamp buf ts = field buf (Timestamp.encode ts)

  (* An option is one field carrying its presence, then the value when present.
     A one-field option lead keeps the framing regular. *)
  let option write buf = function
    | None -> field buf "0"
    | Some v ->
        field buf "1";
        write buf v

  (* A list is a count field followed by that many encoded elements. *)
  let list write buf xs =
    int buf (List.length xs);
    List.iter (write buf) xs

  let contents = Buffer.contents
end

(* --- length-prefixed reader --- *)

module Reader = struct
  type t = { text : string; mutable pos : int }

  exception Bad of string
  exception Unknown of string (* an exercise id not in the catalog *)

  let create text = { text; pos = 0 }

  let field r =
    let len = String.length r.text in
    (* Read the decimal length up to the colon. *)
    let start = r.pos in
    let rec scan i =
      if i >= len then raise (Bad "truncated field length")
      else if r.text.[i] = ':' then i
      else scan (i + 1)
    in
    let colon = scan start in
    let n =
      match int_of_string_opt (String.sub r.text start (colon - start)) with
      | Some n when n >= 0 -> n
      | _ -> raise (Bad "bad field length")
    in
    let value_start = colon + 1 in
    if value_start + n > len then raise (Bad "field overruns input");
    r.pos <- value_start + n;
    String.sub r.text value_start n

  let string = field

  let int r =
    match int_of_string_opt (field r) with
    | Some n -> n
    | None -> raise (Bad "expected int")

  let float r =
    match float_of_string_opt (field r) with
    | Some x -> x
    | None -> raise (Bad "expected float")

  let timestamp r =
    match Timestamp.decode (field r) with
    | Ok ts -> ts
    | Error _ -> raise (Bad "bad timestamp")

  let option read r =
    match field r with
    | "0" -> None
    | "1" -> Some (read r)
    | _ -> raise (Bad "bad option lead")

  let list read r =
    let n = int r in
    if n < 0 then raise (Bad "negative list length");
    let rec go acc i =
      if i = 0 then List.rev acc else go (read r :: acc) (i - 1)
    in
    go [] n

  (* Resolve an exercise id string into an [Exercise.id]. The id type is
     private, so a valid id can only come from the catalog. The reader reports a
     stored id the catalog no longer names as [Unknown]. *)
  let exercise_id r =
    let raw = field r in
    match Exercise.find_id raw with
    | Some e -> Exercise.id e
    | None -> raise (Unknown raw)

  (* Resolve an exercise id string into the catalog entry itself. *)
  let exercise r =
    let raw = field r in
    match Exercise.find_id raw with Some e -> e | None -> raise (Unknown raw)

  let ensure_consumed r =
    if r.pos <> String.length r.text then raise (Bad "trailing bytes")

  (* Run a reader, turning framing faults and unknown ids into typed errors. *)
  let run text f =
    let r = create text in
    match f r with
    | v ->
        ensure_consumed r;
        Ok v
    | exception Bad detail -> Error (Malformed detail)
    | exception Unknown id -> Error (Unknown_exercise id)
end

(* --- observation wire fields ---

   A promoted observation always serializes as the same four fields, in the same
   order: an exercise id, a load, a rep count, and a performed time. The snapshot
   path keeps the id as a raw [Exercise.id]. The record path resolves it to a
   catalog [Exercise.t] and rebuilds the effort at [Positive_failure]. Both share
   this framing, so the byte layout lives here once. *)

module Observation_wire = struct
  (* Write the four fields from already-extracted primitive parts. *)
  let write buf ~exercise_id ~load ~reps ~performed_at =
    Writer.string buf exercise_id;
    Writer.float buf load;
    Writer.int buf reps;
    Writer.timestamp buf performed_at

  (* Read the four fields in order, resolving the leading exercise id with
     [resolve]. The snapshot path resolves to an [Exercise.id]. The record path
     resolves to a catalog [Exercise.t]. The remaining three fields are shared. *)
  let read r ~resolve ~build =
    let exercise = resolve r in
    let load = Reader.float r in
    let reps = Reader.int r in
    let performed_at = Reader.timestamp r in
    build ~exercise ~load ~reps ~performed_at
end

(* --- source set --- *)

let set_type_code : WI.source_set_type -> string = function
  | Normal -> "N"
  | Failure -> "F"
  | Dropset -> "D"

let set_type_of_code = function
  | "N" -> WI.Normal
  | "F" -> WI.Failure
  | "D" -> WI.Dropset
  | other -> raise (Reader.Bad (Printf.sprintf "set type %S" other))

let write_source_set buf (s : WI.source_set) =
  Writer.int buf s.row;
  Writer.string buf s.exercise_name;
  Writer.int buf s.set_index;
  Writer.string buf (set_type_code s.set_type);
  Writer.float buf s.weight_kg;
  Writer.int buf s.reps;
  Writer.option Writer.int buf s.superset_id

let read_source_set r : WI.source_set =
  let row = Reader.int r in
  let exercise_name = Reader.string r in
  let set_index = Reader.int r in
  let set_type = set_type_of_code (Reader.string r) in
  let weight_kg = Reader.float r in
  let reps = Reader.int r in
  let superset_id = Reader.option Reader.int r in
  { row; exercise_name; set_index; set_type; weight_kg; reps; superset_id }

(* --- warning --- *)

let write_warning buf (w : WI.warning) =
  Writer.int buf w.row;
  Writer.option Writer.string buf w.exercise_name;
  match w.kind with
  | Warmup_dropped -> Writer.string buf "warmup"
  | Unsupported_row detail ->
      Writer.string buf "unsupported";
      Writer.string buf detail

let read_warning r : WI.warning =
  let row = Reader.int r in
  let exercise_name = Reader.option Reader.string r in
  let kind =
    match Reader.string r with
    | "warmup" -> WI.Warmup_dropped
    | "unsupported" -> WI.Unsupported_row (Reader.string r)
    | other -> raise (Reader.Bad (Printf.sprintf "warning kind %S" other))
  in
  { row; exercise_name; kind }

(* --- mapping --- *)

let write_mapping buf ((name, id) : string * Exercise.id) =
  Writer.string buf name;
  Writer.string buf (id :> string)

let read_mapping r =
  let name = Reader.string r in
  let id = Reader.exercise_id r in
  (name, id)

(* --- observation snapshot --- *)

let write_observation buf (o : WI.observation_snapshot) =
  Observation_wire.write buf
    ~exercise_id:(o.exercise_id :> string)
    ~load:o.load ~reps:o.reps ~performed_at:o.performed_at

(* Read one observation snapshot, resolving its exercise id to an [Exercise.id]
   through the catalog. *)
let read_observation_snapshot r : WI.observation_snapshot =
  Observation_wire.read r ~resolve:Reader.exercise_id
    ~build:(fun ~exercise ~load ~reps ~performed_at ->
      { WI.exercise_id = exercise; load; reps; performed_at })

(* --- promotion snapshot --- *)

let write_promotion buf (p : WI.promoted_snapshot) =
  Writer.string buf p.title;
  Writer.timestamp buf p.started_at;
  Writer.timestamp buf p.ended_at;
  Writer.list write_observation buf p.observations;
  Writer.list Writer.int buf p.source_rows;
  Writer.option Writer.string buf p.normal_confirmation;
  Writer.int buf p.revision

(* --- workout snapshot --- *)

let write_workout buf (w : WI.workout_snapshot) =
  Writer.string buf (WI.workout_id_to_string w.id);
  Writer.string buf w.title;
  Writer.timestamp buf w.started_at;
  Writer.timestamp buf w.ended_at;
  Writer.list write_source_set buf w.sets;
  Writer.list write_mapping buf w.mappings;
  Writer.option Writer.string buf w.acknowledgement;
  Writer.int buf w.revision;
  Writer.option write_promotion buf w.promotion

(* --- snapshot --- *)

let encode_snapshot (s : WI.snapshot) =
  let buf = Writer.create () in
  Writer.string buf (WI.batch_id_to_string s.batch_id);
  Writer.string buf s.source;
  Writer.string buf s.fingerprint;
  Writer.timestamp buf s.imported_at;
  Writer.list write_workout buf s.workouts;
  Writer.list write_warning buf s.warnings;
  Writer.contents buf

(* The reader builds the parser-free {!WI.snapshot}, keeping every exercise id
   as a raw string. Restoration through {!WI.restore} resolves the ids and
   rebuilds the opaque batch. The use case therefore reports an unknown id or an
   invalid time once. *)
let read_promotion r : WI.promoted_snapshot =
  let title = Reader.string r in
  let started_at = Reader.timestamp r in
  let ended_at = Reader.timestamp r in
  let observations = Reader.list read_observation_snapshot r in
  let source_rows = Reader.list Reader.int r in
  let normal_confirmation = Reader.option Reader.string r in
  let revision = Reader.int r in
  {
    WI.title;
    started_at;
    ended_at;
    observations;
    source_rows;
    normal_confirmation;
    revision;
  }

let read_workout r : WI.workout_snapshot =
  let id = WI.workout_id (Reader.string r) in
  let title = Reader.string r in
  let started_at = Reader.timestamp r in
  let ended_at = Reader.timestamp r in
  let sets = Reader.list read_source_set r in
  let mappings = Reader.list read_mapping r in
  let acknowledgement = Reader.option Reader.string r in
  let revision = Reader.int r in
  let promotion = Reader.option read_promotion r in
  {
    WI.id;
    title;
    started_at;
    ended_at;
    sets;
    mappings;
    acknowledgement;
    revision;
    promotion;
  }

let read_snapshot r : WI.snapshot =
  let batch_id = WI.batch_id (Reader.string r) in
  let source = Reader.string r in
  let fingerprint = Reader.string r in
  let imported_at = Reader.timestamp r in
  let workouts = Reader.list read_workout r in
  let warnings = Reader.list read_warning r in
  { WI.batch_id; source; fingerprint; imported_at; workouts; warnings }

let restore_error : WI.restore_error -> error = function
  | WI.Unknown_exercise id -> Unknown_exercise id
  | WI.Invalid_time -> Invalid_time
  | WI.Invalid_effort e ->
      Malformed
        (Format.asprintf "observation %a" Evidence.Stimulus.Effort.pp_error e)

let decode_snapshot ~find_exercise text =
  match Reader.run text read_snapshot with
  | Error _ as e -> e
  | Ok snapshot -> (
      match WI.restore ~find_exercise snapshot with
      | Ok batch -> Ok batch
      | Error e -> Error (restore_error e))

(* --- imported record --- *)

let encode_record (rec_ : Repository.imported_record) =
  let buf = Writer.create () in
  Writer.string buf (WI.batch_id_to_string rec_.batch_id);
  Writer.string buf (WI.workout_id_to_string rec_.workout_id);
  Writer.int buf rec_.revision;
  Writer.string buf rec_.title;
  Writer.timestamp buf rec_.started_at;
  Writer.timestamp buf rec_.ended_at;
  (* An observation is exercise id, load, reps, performed time: a promoted
     outcome is always Positive_failure, so the outcome is implicit. *)
  Writer.list
    (fun buf obs ->
      let effort = Evidence.Observation.effort obs in
      Observation_wire.write buf
        ~exercise_id:
          (Exercise.id (Evidence.Stimulus.Effort.exercise effort) :> string)
        ~load:(Evidence.Stimulus.Effort.load effort)
        ~reps:(Evidence.Stimulus.Effort.reps effort)
        ~performed_at:(Evidence.Observation.performed_at obs))
    buf rec_.observations;
  (* Provenance. The batch and workout ids repeat the record's own; storing them
     keeps the provenance self-contained after a decode. *)
  Writer.string buf (WI.batch_id_to_string rec_.provenance.batch_id);
  Writer.string buf (WI.workout_id_to_string rec_.provenance.workout_id);
  Writer.list Writer.int buf rec_.provenance.source_rows;
  Writer.option Writer.string buf rec_.provenance.normal_confirmation;
  Writer.int buf rec_.provenance.revision;
  Writer.contents buf

(* Read the record's framing. Observations resolve their exercise through the
   catalog during the read, so an unknown id surfaces as [Unknown_exercise]. *)
let read_record r : Repository.imported_record =
  let batch_id = Reader.string r in
  let workout_id = Reader.string r in
  let revision = Reader.int r in
  let title = Reader.string r in
  let started_at = Reader.timestamp r in
  let ended_at = Reader.timestamp r in
  let observations =
    Reader.list
      (fun r ->
        Observation_wire.read r ~resolve:Reader.exercise
          ~build:(fun ~exercise ~load ~reps ~performed_at ->
            let effort =
              match
                Evidence.Stimulus.Effort.make ~exercise ~load ~reps
                  ~outcome:Evidence.Stimulus.Positive_failure
              with
              | Ok effort -> effort
              | Error e ->
                  raise
                    (Reader.Bad
                       (Format.asprintf "observation %a"
                          Evidence.Stimulus.Effort.pp_error e))
            in
            Evidence.Observation.make ~exercise ~effort ~performed_at))
      r
  in
  let prov_batch_id = Reader.string r in
  let prov_workout_id = Reader.string r in
  let source_rows = Reader.list Reader.int r in
  let normal_confirmation = Reader.option Reader.string r in
  let prov_revision = Reader.int r in
  {
    Repository.batch_id = WI.batch_id batch_id;
    workout_id = WI.workout_id workout_id;
    revision;
    title;
    started_at;
    ended_at;
    observations;
    provenance =
      {
        WI.batch_id = WI.batch_id prov_batch_id;
        workout_id = WI.workout_id prov_workout_id;
        source_rows;
        normal_confirmation;
        revision = prov_revision;
      };
  }

let decode_record text =
  match Reader.run text read_record with
  | Error _ as e -> e
  | Ok record ->
      if
        Recovery.timestamp_to_unix_seconds record.Repository.ended_at
        < Recovery.timestamp_to_unix_seconds record.Repository.started_at
      then Error Invalid_time
      else Ok record