open Lwt.Infix
open Hito_app

exception Corrupt of Codec.error
exception Import_corrupt of Import_codec.error

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

(* --- 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 delete_workout =
    (t2 string string ->. unit)
      "DELETE FROM workout 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 = ?"

  (* --- workout import --- *)

  let upsert_import_batch =
    (t6 string string string string string string ->. unit)
      "INSERT INTO import_batch (trainee_id, batch_id, imported_at, source, \
       fingerprint, snapshot) VALUES (?, ?, ?, ?, ?, ?) ON CONFLICT \
       (trainee_id, batch_id) DO UPDATE SET imported_at = \
       excluded.imported_at, source = excluded.source, fingerprint = \
       excluded.fingerprint, snapshot = excluded.snapshot"

  let import_batch_by_id =
    (t2 string string ->? string)
      "SELECT snapshot FROM import_batch WHERE trainee_id = ? AND batch_id = ?"

  let import_batches =
    (string ->* string)
      "SELECT snapshot FROM import_batch WHERE trainee_id = ? ORDER BY \
       imported_at DESC, batch_id DESC"

  let delete_import_batch =
    (t2 string string ->. unit)
      "DELETE FROM import_batch WHERE trainee_id = ? AND batch_id = ?"

  let import_batch_exists =
    (t2 string string ->? int)
      "SELECT 1 FROM import_batch WHERE trainee_id = ? AND batch_id = ?"

  let upsert_imported_record =
    (t2 (t4 string string string int) (t4 string string string string) ->. unit)
      "INSERT INTO imported_record (trainee_id, workout_id, batch_id, \
       revision, title, started_at, ended_at, snapshot) VALUES (?, ?, ?, ?, ?, \
       ?, ?, ?) ON CONFLICT (trainee_id, workout_id) DO UPDATE SET batch_id = \
       excluded.batch_id, revision = excluded.revision, title = \
       excluded.title, started_at = excluded.started_at, ended_at = \
       excluded.ended_at, snapshot = excluded.snapshot"

  let imported_records =
    (string ->* string)
      "SELECT snapshot FROM imported_record WHERE trainee_id = ? ORDER BY \
       ended_at DESC"
end

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

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

let q_exec t query data =
  run t (fun (module Db : Caqti_lwt.CONNECTION) -> Db.exec query data)

let q_find t query data =
  run t (fun (module Db : Caqti_lwt.CONNECTION) -> Db.find query data)

let q_find_opt t query data =
  run t (fun (module Db : Caqti_lwt.CONNECTION) -> Db.find_opt query data)

let q_collect_list t query data =
  run t (fun (module Db : Caqti_lwt.CONNECTION) -> Db.collect_list query data)

(* SQLite admits one writer at a time. Without a busy timeout a second writer
   fails at once with SQLITE_BUSY; with one it waits its turn. The driver reads
   [busy_timeout] (milliseconds) from the URI, so a caller who sets it keeps
   their value and everyone else gets a five-second wait, which serialises
   overlapping transactions instead of failing them. *)
let default_busy_timeout_ms = 5_000

let with_busy_timeout uri =
  match Uri.get_query_param uri "busy_timeout" with
  | Some _ -> uri
  | None ->
      Uri.add_query_param' uri
        ("busy_timeout", string_of_int default_busy_timeout_ms)

let connect ?pool_config uri =
  match
    Caqti_lwt_unix.connect_pool ?pool_config
      (with_busy_timeout (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
      Lwt_result.return { pool }

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

let find_routine_by_name = Catalog.find_by_name

(* A stored row the codec refuses is corrupt: a storage invariant failed, not a
   user error. A targeted read of one record raises [Corrupt], because the caller
   asked for that record and silently answering "not found" would lose it. A
   list read quarantines the row instead: it is skipped and reported here, so
   one bad row never hides a trainee's whole history. *)
let warn_corrupt ~what ~id detail =
  Printf.eprintf "hito: corrupt %s %S skipped: %s\n%!" what id detail

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

(* Decode one history row, or quarantine it. *)
let decode_history_row (wid, encoded) =
  match Codec.decode_workout ~find_routine:find_routine_by_name encoded with
  | Ok workout -> Some { Repository.id = Repository.workout_id wid; workout }
  | Error e ->
      warn_corrupt ~what:"workout" ~id:wid
        (Format.asprintf "%a" Codec.pp_error e);
      None

(* --- accounts --- *)

(* Create the account in one transaction: check the name is free, take the next
   id, and insert, all under SQLite's write lock. The id comes from the table
   itself on every call, never from a process-local counter, so two concurrent
   registrations cannot mint the same id and a file edited outside this process
   cannot drift from an in-memory count. The UNIQUE constraint on username is
   the final arbiter; the transaction makes the check-then-insert atomic. *)
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) ->
      let open Lwt_result.Syntax in
      Db.with_transaction (fun () ->
          let* existing = Db.find_opt Q.trainee_by_username username_s in
          match existing with
          | Some _ -> Lwt_result.return (Error Repository.Username_taken)
          | None ->
              let* next = Db.find Q.next_trainee () in
              let id = Printf.sprintf "t%d" next in
              let* () =
                Db.exec Q.insert_trainee (id, username_s, credential_s)
              in
              Lwt_result.return
                (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 =
  q_find_opt t Q.trainee_by_username (Trainee.username_to_string username)
  >|= function
  | Some row -> trainee_of_row row
  | None -> None

let find_trainee t id =
  q_find_opt t 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. Both statements run on the same connection, so the check and the
   write cannot race 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 Repository.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 Repository.Username_taken)))

let update_credential t id credential =
  let trainee = Trainee.id_to_string id in
  q_exec t 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 =
  q_find_opt t Q.get_active (Trainee.id_to_string id)
  >|= Option.map Repository.routine_id

let set_active_routine t id (routine : Repository.routine_id) =
  q_exec t Q.set_active (Trainee.id_to_string id, (routine :> string))

let in_progress t id =
  q_find_opt t 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 -> q_exec t Q.clear_in_progress trainee
  | Some workout -> q_exec t 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) =
  q_find_opt t 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
  q_find_opt t Q.workout_by_id (trainee, (record.Repository.id :> string))
  >>= function
  | None -> Lwt.return false
  | Some _ ->
      q_exec t Q.update_workout
        ( encode record.Repository.workout,
          trainee,
          (record.Repository.id :> string) )
      >|= fun () -> true

let delete_workout t id (wid : Repository.workout_id) =
  let trainee = Trainee.id_to_string id in
  q_find_opt t Q.workout_by_id (trainee, (wid :> string)) >>= function
  | None -> Lwt.return false
  | Some _ ->
      q_exec t Q.delete_workout (trainee, (wid :> string)) >|= fun () -> true

let history t id =
  q_collect_list t Q.history (Trainee.id_to_string id)
  >|= List.filter_map decode_history_row

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
  q_find t Q.next_feedback_seq trainee >>= fun seq ->
  q_exec t Q.insert_feedback (trainee, seq, encoded)

(* Quarantine a corrupt row in a list read, naming its position. *)
let quarantine ~what ~pp_error decode rows =
  List.mapi (fun i encoded -> (i + 1, decode encoded)) rows
  |> List.filter_map (fun (position, decoded) ->
      match decoded with
      | Ok v -> Some v
      | Error e ->
          warn_corrupt ~what ~id:(string_of_int position)
            (Format.asprintf "%a" pp_error e);
          None)

let feedback t id =
  q_collect_list t Q.feedback (Trainee.id_to_string id)
  >|= quarantine ~what:"feedback report" ~pp_error:Codec.pp_error
        Codec.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 ->
      q_find t Q.next_app_feedback_seq trainee >>= fun seq ->
      q_exec t 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 =
  q_collect_list t Q.app_feedback () >>= fun rows ->
  q_collect_list t 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
      q_find_opt t Q.feedback_owner (trainee, seq) >>= function
      | None -> Lwt.return false
      | Some owner when String.equal owner voter -> Lwt.return false
      | Some _ -> (
          q_find_opt t Q.vote_exists (trainee, seq, voter) >>= function
          | Some _ -> Lwt.return false
          | None ->
              q_exec t 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) -> (
      q_find_opt t Q.feedback_owner (trainee, seq) >>= function
      | Some owner when String.equal owner (Trainee.id_to_string author) ->
          q_exec t 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) -> (
      q_find_opt t Q.feedback_owner (trainee, seq) >>= function
      | Some owner when String.equal owner (Trainee.id_to_string author) ->
          q_exec t Q.delete_feedback_votes (trainee, seq) >>= fun () ->
          q_exec t Q.delete_app_feedback (trainee, seq) >|= fun () -> true
      | _ -> Lwt.return false)

(* --- workout import --- *)

(* Snapshots resolve their exercise ids through the catalog. The id type is
   private, so the lookup goes through the raw string. *)
let find_exercise (id : Exercise.id) = Exercise.find_id (id :> string)

let decode_batch encoded =
  match Import_codec.decode_snapshot ~find_exercise encoded with
  | Ok batch -> batch
  | Error e -> raise (Import_corrupt e)

let save_import_batch t id batch =
  let trainee = Trainee.id_to_string id in
  let batch_id =
    Workout_import.batch_id_to_string (Workout_import.batch_identity batch)
  in
  let imported_at = Timestamp.encode (Workout_import.batch_imported_at batch) in
  let source = Workout_import.batch_source batch in
  let fingerprint = Workout_import.batch_fingerprint batch in
  let snapshot = Import_codec.encode_snapshot (Workout_import.snapshot batch) in
  q_exec t Q.upsert_import_batch
    (trainee, batch_id, imported_at, source, fingerprint, snapshot)

let find_import_batch t id batch_id =
  q_find_opt t Q.import_batch_by_id
    (Trainee.id_to_string id, Workout_import.batch_id_to_string batch_id)
  >|= function
  | None -> None
  | Some encoded -> Some (decode_batch encoded)

let import_batches t id =
  q_collect_list t Q.import_batches (Trainee.id_to_string id)
  >|= quarantine ~what:"import batch" ~pp_error:Import_codec.pp_error
        (Import_codec.decode_snapshot ~find_exercise)

(* Deleting a provisional batch leaves its promoted records: there is no foreign
   key, so no cascade can reach them. *)
let delete_import_batch t id batch_id =
  let trainee = Trainee.id_to_string id in
  let batch_id = Workout_import.batch_id_to_string batch_id in
  q_find_opt t Q.import_batch_exists (trainee, batch_id) >>= function
  | None -> Lwt.return false
  | Some _ ->
      q_exec t Q.delete_import_batch (trainee, batch_id) >|= fun () -> true

let replace_imported_record t id (record : Repository.imported_record) =
  let trainee = Trainee.id_to_string id in
  let workout_id = Workout_import.workout_id_to_string record.workout_id in
  let batch_id = Workout_import.batch_id_to_string record.batch_id in
  let title = record.title in
  let started_at = Timestamp.encode record.started_at in
  let ended_at = Timestamp.encode record.ended_at in
  let snapshot = Import_codec.encode_record record in
  q_exec t Q.upsert_imported_record
    ( (trainee, workout_id, batch_id, record.revision),
      (title, started_at, ended_at, snapshot) )

let imported_records t id =
  q_collect_list t Q.imported_records (Trainee.id_to_string id)
  >|= quarantine ~what:"imported record" ~pp_error:Import_codec.pp_error
        Import_codec.decode_record