open Lwt.Infix

exception Corrupt of Codec.error

type t = {
  pool : (Caqti_lwt.connection, Caqti_error.t) Caqti_lwt_unix.Pool.t;
  mutable next_trainee : int;
}

(* --- request definitions --- *)

module Q = struct
  open Caqti_request.Infix
  open Caqti_type.Std

  let insert_trainee =
    (t3 string string string ->. unit)
      "INSERT INTO trainee (id, username, credential) VALUES (?, ?, ?)"

  let trainee_by_username =
    (string ->? t3 string string string)
      "SELECT id, username, credential FROM trainee WHERE username = ?"

  let trainee_by_id =
    (string ->? t3 string string string)
      "SELECT id, username, credential FROM trainee WHERE id = ?"

  let next_trainee =
    (unit ->! int)
      "SELECT COALESCE(MAX(CAST(SUBSTR(id, 2) AS INTEGER)), 0) + 1 FROM \
       trainee WHERE id GLOB 't[0-9]*'"

  let update_username =
    (t2 string string ->. unit) "UPDATE trainee SET username = ? WHERE id = ?"

  let update_credential =
    (t2 string string ->. unit) "UPDATE trainee SET credential = ? WHERE id = ?"

  let get_active =
    (string ->? string)
      "SELECT routine_id FROM active_routine WHERE trainee_id = ?"

  let set_active =
    (t2 string string ->. unit)
      "INSERT INTO active_routine (trainee_id, routine_id) VALUES (?, ?) ON \
       CONFLICT (trainee_id) DO UPDATE SET routine_id = excluded.routine_id"

  let get_in_progress =
    (string ->? string) "SELECT encoded FROM in_progress WHERE trainee_id = ?"

  let set_in_progress =
    (t2 string string ->. unit)
      "INSERT INTO in_progress (trainee_id, encoded) VALUES (?, ?) ON CONFLICT \
       (trainee_id) DO UPDATE SET encoded = excluded.encoded"

  let clear_in_progress =
    (string ->. unit) "DELETE FROM in_progress WHERE trainee_id = ?"

  let next_seq =
    (string ->! int)
      "SELECT COALESCE(MAX(seq), 0) + 1 FROM workout WHERE trainee_id = ?"

  let insert_workout =
    (t4 string string int string ->. unit)
      "INSERT INTO workout (id, trainee_id, seq, encoded) VALUES (?, ?, ?, ?)"

  let workout_by_id =
    (t2 string string ->? string)
      "SELECT encoded FROM workout WHERE trainee_id = ? AND id = ?"

  let update_workout =
    (t3 string string string ->. unit)
      "UPDATE workout SET encoded = ? WHERE trainee_id = ? AND id = ?"

  let history =
    (string ->* t2 string string)
      "SELECT id, encoded FROM workout WHERE trainee_id = ? ORDER BY seq DESC"

  let next_feedback_seq =
    (string ->! int)
      "SELECT COALESCE(MAX(seq), 0) + 1 FROM feedback WHERE trainee_id = ?"

  let insert_feedback =
    (t3 string int string ->. unit)
      "INSERT INTO feedback (trainee_id, seq, encoded) VALUES (?, ?, ?)"

  let feedback =
    (string ->* string)
      "SELECT encoded FROM feedback WHERE trainee_id = ? ORDER BY seq DESC"

  let next_app_feedback_seq =
    (string ->! int)
      "SELECT COALESCE(MAX(seq), 0) + 1 FROM app_feedback WHERE trainee_id = ?"

  let insert_app_feedback =
    (t4 string int string string ->. unit)
      "INSERT INTO app_feedback (trainee_id, seq, submitted_at, message) \
       VALUES (?, ?, ?, ?)"

  let app_feedback =
    (unit ->* t7 string int string string string int int)
      "SELECT f.trainee_id, f.seq, t.username, f.submitted_at, f.message, \
       (SELECT COUNT(*) FROM app_feedback c WHERE c.trainee_id = \
       f.trainee_id), (SELECT COUNT(*) FROM app_feedback_vote v WHERE \
       v.feedback_trainee_id = f.trainee_id AND v.feedback_seq = f.seq) FROM \
       app_feedback f JOIN trainee t ON t.id = f.trainee_id ORDER BY 7 DESC, \
       f.submitted_at DESC, f.trainee_id, f.seq DESC"

  let viewer_votes =
    (string ->* t2 string int)
      "SELECT feedback_trainee_id, feedback_seq FROM app_feedback_vote WHERE \
       voter_id = ?"

  let feedback_owner =
    (t2 string int ->? string)
      "SELECT trainee_id FROM app_feedback WHERE trainee_id = ? AND seq = ?"

  let vote_exists =
    (t3 string int string ->? int)
      "SELECT 1 FROM app_feedback_vote WHERE feedback_trainee_id = ? AND \
       feedback_seq = ? AND voter_id = ?"

  let insert_vote =
    (t3 string int string ->. unit)
      "INSERT INTO app_feedback_vote (feedback_trainee_id, feedback_seq, \
       voter_id) VALUES (?, ?, ?)"

  let update_app_feedback =
    (t3 string string int ->. unit)
      "UPDATE app_feedback SET message = ? WHERE trainee_id = ? AND seq = ?"

  let delete_feedback_votes =
    (t2 string int ->. unit)
      "DELETE FROM app_feedback_vote WHERE feedback_trainee_id = ? AND \
       feedback_seq = ?"

  let delete_app_feedback =
    (t2 string int ->. unit)
      "DELETE FROM app_feedback WHERE trainee_id = ? AND seq = ?"
end

(* --- pool helper --- *)

let run t f = Caqti_lwt_unix.Pool.use f t.pool >>= Caqti_lwt.or_fail

let connect uri =
  match Caqti_lwt_unix.connect_pool (Uri.of_string uri) with
  | Error e -> Lwt.return (Error (e :> Caqti_error.t))
  | Ok pool ->
      let open Lwt_result.Syntax in
      let* () = Caqti_lwt_unix.Pool.use Migrations.apply pool in
      let* next_trainee =
        Caqti_lwt_unix.Pool.use
          (fun (module Db : Caqti_lwt.CONNECTION) -> Db.find Q.next_trainee ())
          pool
      in
      Lwt_result.return { pool; next_trainee }

(* --- decoding stored workouts --- *)

let find_routine_by_name = Catalog.find_by_name

let decode encoded =
  match Codec.decode_workout ~find_routine:find_routine_by_name encoded with
  | Ok workout -> workout
  | Error e -> raise (Corrupt e)

(* --- accounts --- *)

let create_trainee t ~username ~credential =
  let username_s = Trainee.username_to_string username in
  let credential_s = Trainee.credential_to_hash credential in
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      Db.find_opt Q.trainee_by_username username_s)
  >>= function
  | Some _ -> Lwt.return (Error `Username_taken)
  | None ->
      let id = Printf.sprintf "t%d" t.next_trainee in
      t.next_trainee <- t.next_trainee + 1;
      run t (fun (module Db : Caqti_lwt.CONNECTION) ->
          Db.exec Q.insert_trainee (id, username_s, credential_s))
      >|= fun () -> Ok { Trainee.id = Trainee.id id; username; credential }

let trainee_of_row (id, username, credential) =
  match Trainee.username username with
  | Ok username ->
      Some
        {
          Trainee.id = Trainee.id id;
          username;
          credential = Trainee.credential_of_hash credential;
        }
  | Error _ -> None

let find_trainee_by_username t username =
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      Db.find_opt Q.trainee_by_username (Trainee.username_to_string username))
  >|= function
  | Some row -> trainee_of_row row
  | None -> None

let find_trainee t id =
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      Db.find_opt Q.trainee_by_id (Trainee.id_to_string id))
  >|= function
  | Some row -> trainee_of_row row
  | None -> None

(* Rename in one transaction: check the name is free for anyone else, then
   update. Running both on the same connection keeps the check and the write
   from racing another rename. The trainee's own current name is a no-op. *)
let update_username t id (username : Trainee.username) =
  let trainee = Trainee.id_to_string id in
  let username_s = Trainee.username_to_string username in
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      let open Lwt_result.Syntax in
      Db.with_transaction (fun () ->
          let* holder = Db.find_opt Q.trainee_by_username username_s in
          match holder with
          | Some (holder_id, _, _) when not (String.equal holder_id trainee) ->
              Lwt_result.return (Error `Username_taken)
          | _ ->
              let* () = Db.exec Q.update_username (username_s, trainee) in
              let* row = Db.find_opt Q.trainee_by_id trainee in
              Lwt_result.return
                (match Option.bind row trainee_of_row with
                | Some updated -> Ok updated
                | None -> Error `Username_taken)))

let update_credential t id credential =
  let trainee = Trainee.id_to_string id in
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      Db.exec Q.update_credential
        (Trainee.credential_to_hash credential, trainee))
  >>= fun () -> find_trainee t id

(* --- catalog --- *)

let list_routines _ = Catalog.routines
let find_routine _ id = Catalog.find id

(* The routine name a workout's prescription belongs to, for encoding. Only the
   ideal routine exists, but resolve it honestly rather than hard-code. *)
let routine_name_of workout =
  let workout_name =
    Prescription.Workout.name (Evidence.Workout.prescription workout)
  in
  List.find_map
    (fun (_, routine) ->
      if
        List.exists
          (fun w -> String.equal (Prescription.Workout.name w) workout_name)
          (Prescription.Routine.workouts routine)
      then Some (Prescription.Routine.name routine)
      else None)
    Catalog.routines
  |> Option.value ~default:""

let encode workout =
  Codec.encode_workout ~routine_name:(routine_name_of workout) workout

(* --- per-trainee selection and workout in progress --- *)

let active_routine t id =
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      Db.find_opt Q.get_active (Trainee.id_to_string id))
  >|= Option.map Repository.routine_id

let set_active_routine t id (routine : Repository.routine_id) =
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      Db.exec Q.set_active (Trainee.id_to_string id, (routine :> string)))

let in_progress t id =
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      Db.find_opt Q.get_in_progress (Trainee.id_to_string id))
  >|= function
  | None -> None
  | Some encoded -> Some (decode encoded)

let set_in_progress t id workout =
  let trainee = Trainee.id_to_string id in
  match workout with
  | None ->
      run t (fun (module Db : Caqti_lwt.CONNECTION) ->
          Db.exec Q.clear_in_progress trainee)
  | Some workout ->
      run t (fun (module Db : Caqti_lwt.CONNECTION) ->
          Db.exec Q.set_in_progress (trainee, encode workout))

(* --- history --- *)

(* Save a finished workout and clear the in-progress slot in one transaction,
   so the two rows never disagree after a crash. Both statements run on the same
   connection, checked out once, inside [with_transaction]. *)
let finish_workout t id workout =
  let trainee = Trainee.id_to_string id in
  let encoded = encode workout in
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      let open Lwt_result.Syntax in
      Db.with_transaction (fun () ->
          let* seq = Db.find Q.next_seq trainee in
          let wid = Printf.sprintf "%s:w%d" trainee seq in
          let* () = Db.exec Q.insert_workout (wid, trainee, seq, encoded) in
          let* () = Db.exec Q.clear_in_progress trainee in
          Lwt_result.return
            { Repository.id = Repository.workout_id wid; workout }))

let find t id (wid : Repository.workout_id) =
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      Db.find_opt Q.workout_by_id (Trainee.id_to_string id, (wid :> string)))
  >|= function
  | None -> None
  | Some encoded -> Some { Repository.id = wid; workout = decode encoded }

let replace t id record =
  let trainee = Trainee.id_to_string id in
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      Db.find_opt Q.workout_by_id (trainee, (record.Repository.id :> string)))
  >>= function
  | None -> Lwt.return false
  | Some _ ->
      run t (fun (module Db : Caqti_lwt.CONNECTION) ->
          Db.exec Q.update_workout
            ( encode record.Repository.workout,
              trainee,
              (record.Repository.id :> string) ))
      >|= fun () -> true

let history t id =
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      Db.collect_list Q.history (Trainee.id_to_string id))
  >|= fun rows ->
  List.map
    (fun (wid, encoded) ->
      { Repository.id = Repository.workout_id wid; workout = decode encoded })
    rows

let log t id =
  history t id >|= fun records ->
  List.fold_left
    (fun log record -> Evidence.Log.add log record.Repository.workout)
    Evidence.Log.empty (List.rev records)

(* --- subjective feedback --- *)

let save_feedback t id report =
  let trainee = Trainee.id_to_string id in
  let encoded = Codec.encode_feedback report in
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      Db.find Q.next_feedback_seq trainee)
  >>= fun seq ->
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      Db.exec Q.insert_feedback (trainee, seq, encoded))

let decode_feedback encoded =
  match Codec.decode_feedback encoded with
  | Ok report -> report
  | Error e -> raise (Corrupt e)

let feedback t id =
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      Db.collect_list Q.feedback (Trainee.id_to_string id))
  >|= List.map decode_feedback

let save_app_feedback t id ~submitted_at ~message =
  let trainee = Trainee.id_to_string id in
  find_trainee t id >>= function
  | None ->
      Lwt.fail_with "cannot save application feedback for an unknown trainee"
  | Some author ->
      run t (fun (module Db : Caqti_lwt.CONNECTION) ->
          Db.find Q.next_app_feedback_seq trainee)
      >>= fun seq ->
      run t (fun (module Db : Caqti_lwt.CONNECTION) ->
          Db.exec Q.insert_app_feedback
            (trainee, seq, Timestamp.encode submitted_at, message))
      >|= fun () ->
      Repository.
        {
          feedback_id = app_feedback_id (Printf.sprintf "%s:%d" trainee seq);
          author = Trainee.username_to_string author.Trainee.username;
          contributions = seq;
          submitted_at;
          message;
          upvotes = 0;
          viewer_upvoted = false;
          viewer_owns = true;
        }

let app_feedback t ~viewer =
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      Db.collect_list Q.app_feedback ())
  >>= fun rows ->
  run t (fun (module Db : Caqti_lwt.CONNECTION) ->
      Db.collect_list Q.viewer_votes (Trainee.id_to_string viewer))
  >|= fun voted ->
  List.map
    (fun (trainee, seq, author, submitted_at, message, contributions, upvotes)
       ->
      let submitted_at =
        match Timestamp.decode submitted_at with
        | Ok timestamp -> timestamp
        | Error _ ->
            raise (Corrupt (Codec.Malformed "application feedback timestamp"))
      in
      Repository.
        {
          feedback_id = app_feedback_id (Printf.sprintf "%s:%d" trainee seq);
          author;
          contributions;
          submitted_at;
          message;
          upvotes;
          viewer_upvoted = List.mem (trainee, seq) voted;
          viewer_owns = String.equal trainee (Trainee.id_to_string viewer);
        })
    rows

let split_app_feedback_id id =
  let raw = Repository.app_feedback_id_to_string id in
  match String.rindex_opt raw ':' with
  | None -> None
  | Some separator ->
      let trainee = String.sub raw 0 separator in
      let seq_start = separator + 1 in
      let seq_length = String.length raw - seq_start in
      Option.map
        (fun seq -> (trainee, seq))
        (int_of_string_opt (String.sub raw seq_start seq_length))

let upvote_app_feedback t ~voter id =
  match split_app_feedback_id id with
  | None -> Lwt.return false
  | Some (trainee, seq) -> (
      let voter = Trainee.id_to_string voter in
      run t (fun (module Db : Caqti_lwt.CONNECTION) ->
          Db.find_opt Q.feedback_owner (trainee, seq))
      >>= function
      | None -> Lwt.return false
      | Some owner when String.equal owner voter -> Lwt.return false
      | Some _ -> (
          run t (fun (module Db : Caqti_lwt.CONNECTION) ->
              Db.find_opt Q.vote_exists (trainee, seq, voter))
          >>= function
          | Some _ -> Lwt.return false
          | None ->
              run t (fun (module Db : Caqti_lwt.CONNECTION) ->
                  Db.exec Q.insert_vote (trainee, seq, voter))
              >|= fun () -> true))

let update_app_feedback t ~author id ~message =
  match split_app_feedback_id id with
  | None -> Lwt.return false
  | Some (trainee, seq) -> (
      run t (fun (module Db : Caqti_lwt.CONNECTION) ->
          Db.find_opt Q.feedback_owner (trainee, seq))
      >>= function
      | Some owner when String.equal owner (Trainee.id_to_string author) ->
          run t (fun (module Db : Caqti_lwt.CONNECTION) ->
              Db.exec Q.update_app_feedback (message, trainee, seq))
          >|= fun () -> true
      | _ -> Lwt.return false)

let delete_app_feedback t ~author id =
  match split_app_feedback_id id with
  | None -> Lwt.return false
  | Some (trainee, seq) -> (
      run t (fun (module Db : Caqti_lwt.CONNECTION) ->
          Db.find_opt Q.feedback_owner (trainee, seq))
      >>= function
      | Some owner when String.equal owner (Trainee.id_to_string author) ->
          run t (fun (module Db : Caqti_lwt.CONNECTION) ->
              Db.exec Q.delete_feedback_votes (trainee, seq))
          >>= fun () ->
          run t (fun (module Db : Caqti_lwt.CONNECTION) ->
              Db.exec Q.delete_app_feedback (trainee, seq))
          >|= fun () -> true
      | _ -> Lwt.return false)