(** Session persistence for Matrix clients.

    Implementation using tomlt for TOML serialization and
    xdge for XDG directory management. *)

module Ed25519 = Mirage_crypto_ec.Ed25519
module X25519 = Mirage_crypto_ec.X25519

(* Helper for URI TOML codec *)
let uri_tomlt : Uri.t Tomlt.t =
  Tomlt.map
    ~dec:Uri.of_string
    ~enc:Uri.to_string
    Tomlt.string

(* Helper for User_id TOML codec *)
let user_id_tomlt : Matrix_proto.Id.User_id.t Tomlt.t =
  Tomlt.map
    ~dec:(fun s ->
      match Matrix_proto.Id.User_id.of_string s with
      | Ok id -> id
      | Error _ -> failwith ("Invalid user_id: " ^ s))
    ~enc:Matrix_proto.Id.User_id.to_string
    Tomlt.string

(* Helper for Device_id TOML codec *)
let device_id_tomlt : Matrix_proto.Id.Device_id.t Tomlt.t =
  Tomlt.map
    ~dec:(fun s ->
      match Matrix_proto.Id.Device_id.of_string s with
      | Ok id -> id
      | Error _ -> failwith ("Invalid device_id: " ^ s))
    ~enc:Matrix_proto.Id.Device_id.to_string
    Tomlt.string

(* Helper for Room_id TOML codec *)
let room_id_tomlt : Matrix_proto.Id.Room_id.t Tomlt.t =
  Tomlt.map
    ~dec:(fun s ->
      match Matrix_proto.Id.Room_id.of_string s with
      | Ok id -> id
      | Error _ -> failwith ("Invalid room_id: " ^ s))
    ~enc:Matrix_proto.Id.Room_id.to_string
    Tomlt.string

(* Ptime codec using UTC *)
let ptime_tomlt : Ptime.t Tomlt.t = Tomlt.ptime ~tz_offset_s:0 ()

module Server = struct
  type t = {
    homeserver : Uri.t;
    user_id : Matrix_proto.Id.User_id.t;
  }

  let tomlt : t Tomlt.t =
    Tomlt.Table.(
      obj (fun homeserver user_id -> { homeserver; user_id })
      |> mem "homeserver" uri_tomlt ~enc:(fun t -> t.homeserver)
      |> mem "user_id" user_id_tomlt ~enc:(fun t -> t.user_id)
      |> finish)
end

module Auth = struct
  type t = {
    access_token : string;
    device_id : Matrix_proto.Id.Device_id.t;
    refresh_token : string option;
  }

  let tomlt : t Tomlt.t =
    Tomlt.Table.(
      obj (fun access_token device_id refresh_token ->
        { access_token; device_id; refresh_token })
      |> mem "access_token" Tomlt.string ~enc:(fun t -> t.access_token)
      |> mem "device_id" device_id_tomlt ~enc:(fun t -> t.device_id)
      |> opt_mem "refresh_token" Tomlt.string ~enc:(fun t -> t.refresh_token)
      |> finish)
end

module Sync_state = struct
  type t = {
    next_batch : string option;
    filter_id : string option;
  }

  let tomlt : t Tomlt.t =
    Tomlt.Table.(
      obj (fun next_batch filter_id -> { next_batch; filter_id })
      |> opt_mem "next_batch" Tomlt.string ~enc:(fun t -> t.next_batch)
      |> opt_mem "filter_id" Tomlt.string ~enc:(fun t -> t.filter_id)
      |> finish)
end

module Metadata = struct
  type t = {
    created_at : Ptime.t;
    last_used_at : Ptime.t;
    client_name : string;
  }

  let tomlt : t Tomlt.t =
    Tomlt.Table.(
      obj (fun created_at last_used_at client_name ->
        { created_at; last_used_at; client_name })
      |> mem "created_at" ptime_tomlt ~enc:(fun t -> t.created_at)
      |> mem "last_used_at" ptime_tomlt ~enc:(fun t -> t.last_used_at)
      |> mem "client_name" Tomlt.string ~enc:(fun t -> t.client_name)
      |> finish)
end

module Session_file = struct
  type t = {
    server : Server.t;
    auth : Auth.t;
    sync : Sync_state.t;
    metadata : Metadata.t;
  }

  let tomlt : t Tomlt.t =
    Tomlt.Table.(
      obj (fun server auth sync metadata ->
        { server; auth; sync; metadata })
      |> mem "server" Server.tomlt ~enc:(fun t -> t.server)
      |> mem "auth" Auth.tomlt ~enc:(fun t -> t.auth)
      |> mem "sync" Sync_state.tomlt ~enc:(fun t -> t.sync)
      |> mem "metadata" Metadata.tomlt ~enc:(fun t -> t.metadata)
      |> finish)
end

module Device_keys = struct
  type t = {
    ed25519_public : string;
    ed25519_private : string;
    curve25519_public : string;
    curve25519_private : string;
    uploaded_at : Ptime.t option;
    algorithms : string list;
  }

  let tomlt : t Tomlt.t =
    Tomlt.Table.(
      obj (fun ed25519_public ed25519_private curve25519_public
               curve25519_private uploaded_at algorithms ->
        { ed25519_public; ed25519_private; curve25519_public;
          curve25519_private; uploaded_at; algorithms })
      |> mem "ed25519_public" Tomlt.string ~enc:(fun t -> t.ed25519_public)
      |> mem "ed25519_private" Tomlt.string ~enc:(fun t -> t.ed25519_private)
      |> mem "curve25519_public" Tomlt.string ~enc:(fun t -> t.curve25519_public)
      |> mem "curve25519_private" Tomlt.string ~enc:(fun t -> t.curve25519_private)
      |> opt_mem "uploaded_at" ptime_tomlt ~enc:(fun t -> t.uploaded_at)
      |> mem "algorithms" (Tomlt.list Tomlt.string) ~dec_absent:[]
          ~enc:(fun t -> t.algorithms)
      |> finish)
end

module One_time_key = struct
  type t = {
    key_id : string;
    public : string;
    private_ : string;
    created_at : Ptime.t;
  }

  let tomlt : t Tomlt.t =
    Tomlt.Table.(
      obj (fun key_id public private_ created_at ->
        { key_id; public; private_; created_at })
      |> mem "key_id" Tomlt.string ~enc:(fun t -> t.key_id)
      |> mem "public" Tomlt.string ~enc:(fun t -> t.public)
      |> mem "private" Tomlt.string ~enc:(fun t -> t.private_)
      |> mem "created_at" ptime_tomlt ~enc:(fun t -> t.created_at)
      |> finish)
end

module One_time_keys_file = struct
  type t = {
    target_count : int;
    last_upload_at : Ptime.t option;
    next_key_id : int;
    keys : One_time_key.t list;
    fallback : One_time_key.t option;
    fallback_used : bool;
  }

  let config_tomlt =
    Tomlt.Table.(
      obj (fun target_count last_upload_at next_key_id ->
        (target_count, last_upload_at, next_key_id))
      |> mem "target_count" Tomlt.int ~dec_absent:50
          ~enc:(fun (tc, _, _) -> tc)
      |> opt_mem "last_upload_at" ptime_tomlt
          ~enc:(fun (_, lu, _) -> lu)
      |> mem "next_key_id" Tomlt.int ~dec_absent:0
          ~enc:(fun (_, _, nk) -> nk)
      |> finish)

  let tomlt : t Tomlt.t =
    Tomlt.Table.(
      obj (fun config keys fallback fallback_used ->
        let target_count, last_upload_at, next_key_id = config in
        { target_count; last_upload_at; next_key_id;
          keys; fallback; fallback_used })
      |> mem "config" config_tomlt
          ~enc:(fun t -> (t.target_count, t.last_upload_at, t.next_key_id))
      |> mem "keys" (Tomlt.list One_time_key.tomlt) ~dec_absent:[]
          ~enc:(fun t -> t.keys)
      |> opt_mem "fallback" One_time_key.tomlt ~enc:(fun t -> t.fallback)
      |> mem "fallback_used" Tomlt.bool ~dec_absent:false
          ~enc:(fun t -> t.fallback_used)
      |> finish)
end

module Olm_session = struct
  type t = {
    their_identity_key : string;
    session_id : string;
    pickle : string;
    created_at : Ptime.t;
    last_used_at : Ptime.t;
  }

  let tomlt : t Tomlt.t =
    Tomlt.Table.(
      obj (fun their_identity_key session_id pickle created_at last_used_at ->
        { their_identity_key; session_id; pickle; created_at; last_used_at })
      |> mem "their_identity_key" Tomlt.string
          ~enc:(fun t -> t.their_identity_key)
      |> mem "session_id" Tomlt.string ~enc:(fun t -> t.session_id)
      |> mem "pickle" Tomlt.string ~enc:(fun t -> t.pickle)
      |> mem "created_at" ptime_tomlt ~enc:(fun t -> t.created_at)
      |> mem "last_used_at" ptime_tomlt ~enc:(fun t -> t.last_used_at)
      |> finish)
end

module Olm_sessions_file = struct
  type t = { sessions : Olm_session.t list }

  let tomlt : t Tomlt.t =
    Tomlt.Table.(
      obj (fun sessions -> { sessions })
      |> mem "sessions" (Tomlt.list Olm_session.tomlt) ~dec_absent:[]
          ~enc:(fun t -> t.sessions)
      |> finish)
end

module Megolm_inbound = struct
  type t = {
    room_id : Matrix_proto.Id.Room_id.t;
    session_id : string;
    sender_key : string;
    signing_key : string;
    pickle : string;
    first_known_index : int;
    created_at : Ptime.t;
  }

  let tomlt : t Tomlt.t =
    Tomlt.Table.(
      obj (fun room_id session_id sender_key signing_key pickle
               first_known_index created_at ->
        { room_id; session_id; sender_key; signing_key; pickle;
          first_known_index; created_at })
      |> mem "room_id" room_id_tomlt ~enc:(fun t -> t.room_id)
      |> mem "session_id" Tomlt.string ~enc:(fun t -> t.session_id)
      |> mem "sender_key" Tomlt.string ~enc:(fun t -> t.sender_key)
      |> mem "signing_key" Tomlt.string ~enc:(fun t -> t.signing_key)
      |> mem "pickle" Tomlt.string ~enc:(fun t -> t.pickle)
      |> mem "first_known_index" Tomlt.int ~enc:(fun t -> t.first_known_index)
      |> mem "created_at" ptime_tomlt ~enc:(fun t -> t.created_at)
      |> finish)
end

module Megolm_inbound_file = struct
  type t = { sessions : Megolm_inbound.t list }

  let tomlt : t Tomlt.t =
    Tomlt.Table.(
      obj (fun sessions -> { sessions })
      |> mem "sessions" (Tomlt.list Megolm_inbound.tomlt) ~dec_absent:[]
          ~enc:(fun t -> t.sessions)
      |> finish)
end

module Shared_with = struct
  type t = {
    user_id : Matrix_proto.Id.User_id.t;
    device_id : Matrix_proto.Id.Device_id.t;
    shared_at : Ptime.t;
  }

  let tomlt : t Tomlt.t =
    Tomlt.Table.(
      obj (fun user_id device_id shared_at ->
        { user_id; device_id; shared_at })
      |> mem "user_id" user_id_tomlt ~enc:(fun t -> t.user_id)
      |> mem "device_id" device_id_tomlt ~enc:(fun t -> t.device_id)
      |> mem "shared_at" ptime_tomlt ~enc:(fun t -> t.shared_at)
      |> finish)
end

module Megolm_outbound = struct
  type t = {
    room_id : Matrix_proto.Id.Room_id.t;
    session_id : string;
    pickle : string;
    message_index : int;
    created_at : Ptime.t;
    message_count : int;
    max_age_ms : int64;
    shared_with : Shared_with.t list;
  }

  let tomlt : t Tomlt.t =
    Tomlt.Table.(
      obj (fun room_id session_id pickle message_index created_at
               message_count max_age_ms shared_with ->
        { room_id; session_id; pickle; message_index; created_at;
          message_count; max_age_ms; shared_with })
      |> mem "room_id" room_id_tomlt ~enc:(fun t -> t.room_id)
      |> mem "session_id" Tomlt.string ~enc:(fun t -> t.session_id)
      |> mem "pickle" Tomlt.string ~enc:(fun t -> t.pickle)
      |> mem "message_index" Tomlt.int ~enc:(fun t -> t.message_index)
      |> mem "created_at" ptime_tomlt ~enc:(fun t -> t.created_at)
      |> mem "message_count" Tomlt.int ~enc:(fun t -> t.message_count)
      |> mem "max_age_ms" Tomlt.int64 ~enc:(fun t -> t.max_age_ms)
      |> mem "shared_with" (Tomlt.list Shared_with.tomlt) ~dec_absent:[]
          ~enc:(fun t -> t.shared_with)
      |> finish)
end

module Megolm_outbound_file = struct
  type t = { sessions : Megolm_outbound.t list }

  let tomlt : t Tomlt.t =
    Tomlt.Table.(
      obj (fun sessions -> { sessions })
      |> mem "sessions" (Tomlt.list Megolm_outbound.tomlt) ~dec_absent:[]
          ~enc:(fun t -> t.sessions)
      |> finish)
end

(* ============================================================ *)
(* Pickle Functions using jsont                                 *)
(* ============================================================ *)

module Pickle = struct
  (* Base64 encoding/decoding - Matrix uses unpadded base64 *)
  let base64_encode s = Base64.encode_string ~pad:false s
  let base64_decode s =
    match Base64.decode ~pad:false s with
    | Ok s -> s
    | Error _ -> failwith "Invalid base64"

  (* Jsont codec for Ed25519 private key *)
  let ed25519_priv_jsont : Ed25519.priv Jsont.t =
    Jsont.map
      ~dec:(fun s ->
        let octets = base64_decode s in
        match Ed25519.priv_of_octets octets with
        | Ok priv -> priv
        | Error _ -> failwith "Invalid Ed25519 private key")
      ~enc:(fun priv -> Ed25519.priv_to_octets priv |> base64_encode)
      Jsont.string

  (* Jsont codec for Ed25519 public key *)
  let ed25519_pub_jsont : Ed25519.pub Jsont.t =
    Jsont.map
      ~dec:(fun s ->
        let octets = base64_decode s in
        match Ed25519.pub_of_octets octets with
        | Ok pub -> pub
        | Error _ -> failwith "Invalid Ed25519 public key")
      ~enc:(fun pub -> Ed25519.pub_to_octets pub |> base64_encode)
      Jsont.string

  (* Jsont codec for X25519 secret *)
  let x25519_secret_jsont : X25519.secret Jsont.t =
    Jsont.map
      ~dec:(fun s ->
        let octets = base64_decode s in
        match X25519.secret_of_octets octets with
        | Ok (secret, _) -> secret
        | Error _ -> failwith "Invalid X25519 secret key")
      ~enc:(fun secret -> X25519.secret_to_octets secret |> base64_encode)
      Jsont.string

  (* Jsont codec for Ptime.t as ISO 8601 string *)
  let ptime_jsont : Ptime.t Jsont.t =
    Jsont.map
      ~dec:(fun s ->
        match Ptime.of_rfc3339 s with
        | Ok (t, _, _) -> t
        | Error _ -> failwith "Invalid RFC3339 timestamp")
      ~enc:(fun t -> Ptime.to_rfc3339 ~tz_offset_s:0 t)
      Jsont.string

  (* ------------------------------------------------------------ *)
  (* Olm.Account pickle                                           *)
  (* ------------------------------------------------------------ *)

  type account_pickle = {
    ed25519_priv : Ed25519.priv;
    ed25519_pub : Ed25519.pub;
    curve25519_secret : X25519.secret;
    curve25519_public : string;
    one_time_keys : (string * string * string) list;
    (* key_id, secret_b64, public_b64 *)
    fallback_key : (string * string * string) option;
    next_key_id : int;
    max_one_time_keys : int;
  }

  let account_pickle_jsont : account_pickle Jsont.t =
    Jsont.Object.(
      map (fun ed25519_priv ed25519_pub curve25519_secret curve25519_public
               one_time_keys fallback_key next_key_id max_one_time_keys ->
        { ed25519_priv; ed25519_pub; curve25519_secret; curve25519_public;
          one_time_keys; fallback_key; next_key_id; max_one_time_keys })
      |> mem "ed25519_priv" ed25519_priv_jsont
      |> mem "ed25519_pub" ed25519_pub_jsont
      |> mem "curve25519_secret" x25519_secret_jsont
      |> mem "curve25519_public" Jsont.string
      |> mem "one_time_keys"
          (Jsont.list (Jsont.list Jsont.string |> Jsont.map
            ~dec:(function
              | [a; b; c] -> (a, b, c)
              | _ -> failwith "Expected 3 elements")
            ~enc:(fun (a, b, c) -> [a; b; c])))
      |> opt_mem "fallback_key"
          (Jsont.list Jsont.string |> Jsont.map
            ~dec:(function
              | [a; b; c] -> (a, b, c)
              | _ -> failwith "Expected 3 elements")
            ~enc:(fun (a, b, c) -> [a; b; c]))
      |> mem "next_key_id" Jsont.int
      |> mem "max_one_time_keys" Jsont.int
      |> finish)

  let pickle_account (account : Olm.Account.t) : string =
    let one_time_keys =
      List.map (fun (key_id, (secret, public)) ->
        (key_id, X25519.secret_to_octets secret |> base64_encode,
         base64_encode public))
        account.one_time_keys
    in
    let fallback_key =
      Option.map (fun (key_id, (secret, public)) ->
        (key_id, X25519.secret_to_octets secret |> base64_encode,
         base64_encode public))
        account.fallback_key
    in
    let pickle = {
      ed25519_priv = account.ed25519_priv;
      ed25519_pub = account.ed25519_pub;
      curve25519_secret = account.curve25519_secret;
      curve25519_public = account.curve25519_public;
      one_time_keys;
      fallback_key;
      next_key_id = account.next_key_id;
      max_one_time_keys = account.max_one_time_keys;
    } in
    match Jsont_bytesrw.encode_string account_pickle_jsont pickle with
    | Ok s -> s
    | Error e -> failwith ("Failed to pickle account: " ^ e)

  let unpickle_account (s : string) : (Olm.Account.t, string) result =
    match Jsont_bytesrw.decode_string account_pickle_jsont s with
    | Error e -> Error ("Failed to unpickle account: " ^ e)
    | Ok pickle ->
        let one_time_keys =
          List.map (fun (key_id, secret_b64, public_b64) ->
            let secret =
              match X25519.secret_of_octets (base64_decode secret_b64) with
              | Ok (s, _) -> s
              | Error _ -> failwith "Invalid X25519 secret"
            in
            let public = base64_decode public_b64 in
            (key_id, (secret, public)))
            pickle.one_time_keys
        in
        let fallback_key =
          Option.map (fun (key_id, secret_b64, public_b64) ->
            let secret =
              match X25519.secret_of_octets (base64_decode secret_b64) with
              | Ok (s, _) -> s
              | Error _ -> failwith "Invalid X25519 secret"
            in
            let public = base64_decode public_b64 in
            (key_id, (secret, public)))
            pickle.fallback_key
        in
        Ok {
          Olm.Account.ed25519_priv = pickle.ed25519_priv;
          ed25519_pub = pickle.ed25519_pub;
          curve25519_secret = pickle.curve25519_secret;
          curve25519_public = pickle.curve25519_public;
          one_time_keys;
          fallback_key;
          next_key_id = pickle.next_key_id;
          max_one_time_keys = pickle.max_one_time_keys;
        }

  (* ------------------------------------------------------------ *)
  (* Olm.Session pickle                                           *)
  (* ------------------------------------------------------------ *)

  type chain_key_pickle = {
    key : string;  (* base64 *)
    index : int;
  }

  let chain_key_pickle_jsont : chain_key_pickle Jsont.t =
    Jsont.Object.(
      map (fun key index -> { key; index })
      |> mem "key" Jsont.string
      |> mem "index" Jsont.int
      |> finish)

  type session_pickle = {
    session_id : string;
    their_identity_key : string;
    their_ratchet_key : string option;
    our_ratchet_secret : string;
    our_ratchet_public : string;
    root_key : string;
    sending_chain : chain_key_pickle option;
    receiving_chains : (string * chain_key_pickle) list;
    skipped_keys : ((string * int) * string) list;
    creation_time : Ptime.t;
  }

  (* Encode (string, int) pair as a JSON object with "key" and "index" *)
  let string_int_pair_jsont : (string * int) Jsont.t =
    Jsont.Object.(
      map (fun k i -> (k, i))
      |> mem "key" Jsont.string
      |> mem "index" Jsont.int
      |> finish)

  let session_pickle_jsont : session_pickle Jsont.t =
    Jsont.Object.(
      map (fun session_id their_identity_key their_ratchet_key
               our_ratchet_secret our_ratchet_public root_key sending_chain
               receiving_chains skipped_keys creation_time ->
        { session_id; their_identity_key; their_ratchet_key;
          our_ratchet_secret; our_ratchet_public; root_key; sending_chain;
          receiving_chains; skipped_keys; creation_time })
      |> mem "session_id" Jsont.string
      |> mem "their_identity_key" Jsont.string
      |> opt_mem "their_ratchet_key" Jsont.string
      |> mem "our_ratchet_secret" Jsont.string
      |> mem "our_ratchet_public" Jsont.string
      |> mem "root_key" Jsont.string
      |> opt_mem "sending_chain" chain_key_pickle_jsont
      |> mem "receiving_chains"
          (Jsont.list (Jsont.Object.(
            map (fun k v -> (k, v))
            |> mem "key" Jsont.string
            |> mem "chain" chain_key_pickle_jsont
            |> finish)))
      |> mem "skipped_keys"
          (Jsont.list (Jsont.Object.(
            map (fun idx_key msg_key -> (idx_key, msg_key))
            |> mem "index_key" string_int_pair_jsont
            |> mem "msg_key" Jsont.string
            |> finish)))
      |> mem "creation_time" ptime_jsont
      |> finish)

  let pickle_session (session : Olm.Session.t) : string =
    let sending_chain = Option.map (fun (ck : Olm.Session.chain_key) ->
      { key = base64_encode ck.key; index = ck.index })
      session.sending_chain
    in
    let receiving_chains =
      List.map (fun (rk, (ck : Olm.Session.chain_key)) ->
        (base64_encode rk, { key = base64_encode ck.key; index = ck.index }))
        session.receiving_chains
    in
    let skipped_keys =
      List.map (fun ((rk, idx), mk) ->
        ((base64_encode rk, idx), base64_encode mk))
        session.skipped_keys
    in
    let pickle = {
      session_id = session.session_id;
      their_identity_key = base64_encode session.their_identity_key;
      their_ratchet_key = Option.map base64_encode session.their_ratchet_key;
      our_ratchet_secret = X25519.secret_to_octets session.our_ratchet_secret
                           |> base64_encode;
      our_ratchet_public = base64_encode session.our_ratchet_public;
      root_key = base64_encode session.root_key;
      sending_chain;
      receiving_chains;
      skipped_keys;
      creation_time = session.creation_time;
    } in
    match Jsont_bytesrw.encode_string session_pickle_jsont pickle with
    | Ok s -> s
    | Error e -> failwith ("Failed to pickle session: " ^ e)

  let unpickle_session (s : string) : (Olm.Session.t, string) result =
    match Jsont_bytesrw.decode_string session_pickle_jsont s with
    | Error e -> Error ("Failed to unpickle session: " ^ e)
    | Ok pickle ->
        let our_ratchet_secret =
          match X25519.secret_of_octets (base64_decode pickle.our_ratchet_secret) with
          | Ok (s, _) -> s
          | Error _ -> failwith "Invalid ratchet secret"
        in
        let sending_chain = Option.map (fun p ->
          { Olm.Session.key = base64_decode p.key; index = p.index })
          pickle.sending_chain
        in
        let receiving_chains =
          List.map (fun (rk_b64, p) ->
            (base64_decode rk_b64,
             { Olm.Session.key = base64_decode p.key; index = p.index }))
            pickle.receiving_chains
        in
        let skipped_keys =
          List.map (fun ((rk_b64, idx), mk_b64) ->
            ((base64_decode rk_b64, idx), base64_decode mk_b64))
            pickle.skipped_keys
        in
        Ok {
          Olm.Session.session_id = pickle.session_id;
          their_identity_key = base64_decode pickle.their_identity_key;
          their_ratchet_key = Option.map base64_decode pickle.their_ratchet_key;
          our_ratchet_secret;
          our_ratchet_public = base64_decode pickle.our_ratchet_public;
          root_key = base64_decode pickle.root_key;
          sending_chain;
          receiving_chains;
          skipped_keys;
          creation_time = pickle.creation_time;
        }

  (* ------------------------------------------------------------ *)
  (* Megolm.Inbound pickle                                        *)
  (* ------------------------------------------------------------ *)

  type megolm_inbound_pickle = {
    session_id : string;
    sender_key : string;
    room_id : string;
    ratchet : string list;  (* 4 x base64 strings *)
    message_index : int;
    received_indices : int list;
    signing_key : string;
    creation_time : Ptime.t;
  }

  let megolm_inbound_pickle_jsont : megolm_inbound_pickle Jsont.t =
    Jsont.Object.(
      map (fun session_id sender_key room_id ratchet message_index
               received_indices signing_key creation_time ->
        { session_id; sender_key; room_id; ratchet; message_index;
          received_indices; signing_key; creation_time })
      |> mem "session_id" Jsont.string
      |> mem "sender_key" Jsont.string
      |> mem "room_id" Jsont.string
      |> mem "ratchet" (Jsont.list Jsont.string)
      |> mem "message_index" Jsont.int
      |> mem "received_indices" (Jsont.list Jsont.int)
      |> mem "signing_key" Jsont.string
      |> mem "creation_time" ptime_jsont
      |> finish)

  let pickle_megolm_inbound (session : Olm.Megolm.Inbound.t) : string =
    let ratchet = Array.to_list session.ratchet |> List.map base64_encode in
    let pickle = {
      session_id = session.session_id;
      sender_key = session.sender_key;
      room_id = session.room_id;
      ratchet;
      message_index = session.message_index;
      received_indices = session.received_indices;
      signing_key = session.signing_key;
      creation_time = session.creation_time;
    } in
    match Jsont_bytesrw.encode_string megolm_inbound_pickle_jsont pickle with
    | Ok s -> s
    | Error e -> failwith ("Failed to pickle megolm inbound: " ^ e)

  let unpickle_megolm_inbound (s : string)
      : (Olm.Megolm.Inbound.t, string) result =
    match Jsont_bytesrw.decode_string megolm_inbound_pickle_jsont s with
    | Error e -> Error ("Failed to unpickle megolm inbound: " ^ e)
    | Ok pickle ->
        let ratchet =
          List.map base64_decode pickle.ratchet |> Array.of_list
        in
        Ok {
          Olm.Megolm.Inbound.session_id = pickle.session_id;
          sender_key = pickle.sender_key;
          room_id = pickle.room_id;
          ratchet;
          message_index = pickle.message_index;
          received_indices = pickle.received_indices;
          signing_key = pickle.signing_key;
          creation_time = pickle.creation_time;
        }

  (* ------------------------------------------------------------ *)
  (* Megolm.Outbound pickle                                       *)
  (* ------------------------------------------------------------ *)

  type megolm_outbound_pickle = {
    session_id : string;
    room_id : string;
    ratchet : string list;
    message_index : int;
    signing_priv : Ed25519.priv;
    signing_pub : Ed25519.pub;
    creation_time : Ptime.t;
    message_count : int;
    max_messages : int;
    max_age_s : int;  (* stored as seconds *)
    shared_with : (string * string) list;  (* user_id, device_id pairs *)
  }

  let megolm_outbound_pickle_jsont : megolm_outbound_pickle Jsont.t =
    Jsont.Object.(
      map (fun session_id room_id ratchet message_index signing_priv
               signing_pub creation_time message_count max_messages max_age_s
               shared_with ->
        { session_id; room_id; ratchet; message_index; signing_priv;
          signing_pub; creation_time; message_count; max_messages; max_age_s;
          shared_with })
      |> mem "session_id" Jsont.string
      |> mem "room_id" Jsont.string
      |> mem "ratchet" (Jsont.list Jsont.string)
      |> mem "message_index" Jsont.int
      |> mem "signing_priv" ed25519_priv_jsont
      |> mem "signing_pub" ed25519_pub_jsont
      |> mem "creation_time" ptime_jsont
      |> mem "message_count" Jsont.int
      |> mem "max_messages" Jsont.int
      |> mem "max_age_s" Jsont.int
      |> mem "shared_with"
          (Jsont.list (Jsont.list Jsont.string |> Jsont.map
            ~dec:(function
              | [a; b] -> (a, b)
              | _ -> failwith "Expected 2 elements")
            ~enc:(fun (a, b) -> [a; b])))
      |> finish)

  let pickle_megolm_outbound (session : Olm.Megolm.Outbound.t) : string =
    let ratchet = Array.to_list session.ratchet |> List.map base64_encode in
    let max_age_s =
      match Ptime.Span.to_int_s session.max_age with
      | Some s -> s
      | None -> 604800  (* default 1 week *)
    in
    let pickle = {
      session_id = session.session_id;
      room_id = session.room_id;
      ratchet;
      message_index = session.message_index;
      signing_priv = session.signing_priv;
      signing_pub = session.signing_pub;
      creation_time = session.creation_time;
      message_count = session.message_count;
      max_messages = session.max_messages;
      max_age_s;
      shared_with = session.shared_with;
    } in
    match Jsont_bytesrw.encode_string megolm_outbound_pickle_jsont pickle with
    | Ok s -> s
    | Error e -> failwith ("Failed to pickle megolm outbound: " ^ e)

  let unpickle_megolm_outbound (s : string)
      : (Olm.Megolm.Outbound.t, string) result =
    match Jsont_bytesrw.decode_string megolm_outbound_pickle_jsont s with
    | Error e -> Error ("Failed to unpickle megolm outbound: " ^ e)
    | Ok pickle ->
        let ratchet =
          List.map base64_decode pickle.ratchet |> Array.of_list
        in
        let max_age = Ptime.Span.of_int_s pickle.max_age_s in
        Ok {
          Olm.Megolm.Outbound.session_id = pickle.session_id;
          room_id = pickle.room_id;
          ratchet;
          message_index = pickle.message_index;
          signing_priv = pickle.signing_priv;
          signing_pub = pickle.signing_pub;
          creation_time = pickle.creation_time;
          message_count = pickle.message_count;
          max_messages = pickle.max_messages;
          max_age;
          shared_with = pickle.shared_with;
        }
end

(* ============================================================ *)
(* Session Store                                                *)
(* ============================================================ *)

module Store = struct
  type t = {
    profile_path : Eio.Fs.dir_ty Eio.Path.t;
  }

  let create ~xdg ~profile =
    let data_dir = Xdge.data_dir xdg in
    let profile_path = Eio.Path.(data_dir / "profiles" / profile) in
    (* Ensure directory exists *)
    Eio.Path.mkdirs ~exists_ok:true ~perm:0o700 profile_path;
    { profile_path }

  let profile_dir t = t.profile_path

  let exists t =
    Eio.Path.is_file Eio.Path.(t.profile_path / "session.toml")

  (* Helper to load a TOML file with a codec *)
  let load_toml (type a) (codec : a Tomlt.t) (path : _ Eio.Path.t)
      : a option =
    if Eio.Path.is_file path then
      match Tomlt_eio.decode_file codec path with
      | Ok v -> Some v
      | Error _ -> None
    else
      None

  (* Helper to save a TOML file with a codec *)
  let save_toml (type a) (codec : a Tomlt.t) (path : _ Eio.Path.t) (value : a)
      : unit =
    Tomlt_eio.encode_file codec value path

  let load_session t =
    load_toml Session_file.tomlt Eio.Path.(t.profile_path / "session.toml")

  let save_session t session =
    save_toml Session_file.tomlt Eio.Path.(t.profile_path / "session.toml")
      session

  let load_device_keys t =
    load_toml Device_keys.tomlt Eio.Path.(t.profile_path / "device.toml")

  let save_device_keys t keys =
    save_toml Device_keys.tomlt Eio.Path.(t.profile_path / "device.toml") keys

  let load_one_time_keys t =
    load_toml One_time_keys_file.tomlt
      Eio.Path.(t.profile_path / "one_time_keys.toml")

  let save_one_time_keys t keys =
    save_toml One_time_keys_file.tomlt
      Eio.Path.(t.profile_path / "one_time_keys.toml") keys

  let load_olm_sessions t =
    load_toml Olm_sessions_file.tomlt
      Eio.Path.(t.profile_path / "olm_sessions.toml")

  let save_olm_sessions t sessions =
    save_toml Olm_sessions_file.tomlt
      Eio.Path.(t.profile_path / "olm_sessions.toml") sessions

  let load_megolm_inbound t =
    load_toml Megolm_inbound_file.tomlt
      Eio.Path.(t.profile_path / "megolm_inbound.toml")

  let save_megolm_inbound t sessions =
    save_toml Megolm_inbound_file.tomlt
      Eio.Path.(t.profile_path / "megolm_inbound.toml") sessions

  let load_megolm_outbound t =
    load_toml Megolm_outbound_file.tomlt
      Eio.Path.(t.profile_path / "megolm_outbound.toml")

  let save_megolm_outbound t sessions =
    save_toml Megolm_outbound_file.tomlt
      Eio.Path.(t.profile_path / "megolm_outbound.toml") sessions

  let clear t =
    let files = [
      "session.toml"; "device.toml"; "one_time_keys.toml";
      "olm_sessions.toml"; "megolm_inbound.toml"; "megolm_outbound.toml"
    ] in
    List.iter (fun filename ->
      let path = Eio.Path.(t.profile_path / filename) in
      if Eio.Path.is_file path then
        Eio.Path.unlink path)
      files
end