(** Olm/Megolm cryptographic session management.

    This module implements the Olm double-ratchet algorithm for encrypted
    to-device messages, and Megolm for encrypted room messages.

    @see <https://spec.matrix.org/v1.11/client-server-api/#end-to-end-encryption> E2EE *)

(** {1 Olm Account} *)

(** Olm account - manages identity keys and one-time keys.

    An Olm account contains:
    - Ed25519 identity key (for signing)
    - Curve25519 identity key (for key exchange)
    - One-time keys (for session establishment) *)
module Account : sig
  (** Account type. *)
  type t = {
    ed25519_priv : Mirage_crypto_ec.Ed25519.priv;
    ed25519_pub : Mirage_crypto_ec.Ed25519.pub;
    curve25519_secret : Mirage_crypto_ec.X25519.secret;
    curve25519_public : string;
    mutable one_time_keys : (string * (Mirage_crypto_ec.X25519.secret * string)) list;
    mutable fallback_key : (string * (Mirage_crypto_ec.X25519.secret * string)) option;
    mutable next_key_id : int;
    max_one_time_keys : int;
  }

  val create : unit -> t
  (** Generate a new Olm account with fresh identity keys. *)

  val ed25519_key : t -> string
  (** Get the Ed25519 identity key as base64. *)

  val curve25519_key : t -> string
  (** Get the Curve25519 identity key as base64. *)

  val identity_keys : t -> string * string
  (** Get the identity keys as a pair (ed25519, curve25519). *)

  val sign : t -> string -> string
  (** Sign a message with the account's Ed25519 key. *)

  val generate_one_time_keys : t -> int -> unit
  (** Generate new one-time keys. *)

  val one_time_keys : t -> (string * string) list
  (** Get one-time keys for upload (key_id -> public_key). *)

  val signed_one_time_keys : t -> (string * string * string) list
  (** Get signed one-time keys for upload. Returns [(key_id, public_key, signature)]. *)

  val mark_keys_as_published : t -> unit
  (** Mark one-time keys as published. *)

  val generate_fallback_key : t -> unit
  (** Generate a fallback key. *)

  val fallback_key : t -> (string * string) option
  (** Get the fallback key if one exists. *)

  val remove_one_time_key : t -> string -> unit
  (** Remove a one-time key by ID (after session creation). *)

  val one_time_keys_count : t -> int
  (** Number of unpublished one-time keys. *)

  val max_one_time_keys : t -> int
  (** Maximum number of one-time keys this account can hold. *)
end

(** {1 Olm Session} *)

(** Olm session state for the double ratchet algorithm. *)
module Session : sig
  type chain_key = { key : string; index : int }
  type root_key = string
  type message_key = string

  (** Session type. *)
  type t = {
    session_id : string;
    their_identity_key : string;
    mutable their_ratchet_key : string option;
    mutable our_ratchet_secret : Mirage_crypto_ec.X25519.secret;
    mutable our_ratchet_public : string;
    mutable root_key : root_key;
    mutable sending_chain : chain_key option;
    mutable receiving_chains : (string * chain_key) list;
    mutable skipped_keys : ((string * int) * message_key) list;
    creation_time : Ptime.t;
  }

  val create_outbound :
    Account.t ->
    their_identity_key:string ->
    their_one_time_key:string ->
    t
  (** Create a new outbound session (when sending first message). *)

  val create_inbound :
    Account.t ->
    their_identity_key:string ->
    their_ephemeral_key:string ->
    one_time_key_id:string ->
    t
  (** Create a new inbound session (when receiving first message). *)

  val session_id : t -> string
  (** Get session ID. *)

  val their_identity_key : t -> string
  (** Get their identity key. *)

  val encrypt : t -> string -> int * string
  (** Encrypt a plaintext message. Returns (message_type, ciphertext). *)

  val decrypt : t -> message_type:int -> ciphertext:string -> (string, string) result
  (** Decrypt a message. *)

  val is_pre_key_message : int -> bool
  (** Check if this is a pre-key message (first message in session). *)
end

(** {1 Megolm Sessions} *)

(** Megolm session for room message encryption.

    Megolm uses a ratchet that only moves forward, making it efficient
    for encrypting many messages to many recipients. *)
module Megolm : sig

  (** Inbound session for decrypting received room messages. *)
  module Inbound : sig
    (** Inbound session type. *)
    type t = {
      session_id : string;
      sender_key : string;
      room_id : string;
      mutable ratchet : string array;
      mutable message_index : int;
      mutable received_indices : int list;
      signing_key : string;
      creation_time : Ptime.t;
    }

    val of_export :
      session_id:string ->
      sender_key:string ->
      room_id:string ->
      ratchet:string array ->
      message_index:int ->
      signing_key:string ->
      t
    (** Create from exported session data. *)

    val from_room_key :
      sender_key:string ->
      room_id:string ->
      session_id:string ->
      session_key:string ->
      signing_key:string ->
      t
    (** Create from room key event (m.room_key). *)

    val session_id : t -> string
    (** Get session ID. *)

    val sender_key : t -> string
    (** Get sender key. *)

    val room_id : t -> string
    (** Get room ID. *)

    val first_known_index : t -> int
    (** Get first known message index. *)

    val decrypt : t -> ciphertext:string -> message_index:int -> (string, string) result
    (** Decrypt a message. *)
  end

  (** Outbound session for encrypting messages to send to a room. *)
  module Outbound : sig
    (** Outbound session type. *)
    type t = {
      session_id : string;
      room_id : string;
      mutable ratchet : string array;
      mutable message_index : int;
      signing_priv : Mirage_crypto_ec.Ed25519.priv;
      signing_pub : Mirage_crypto_ec.Ed25519.pub;
      creation_time : Ptime.t;
      mutable message_count : int;
      max_messages : int;
      max_age : Ptime.Span.t;
      mutable shared_with : (string * string) list;
    }

    val create : room_id:string -> t
    (** Create a new outbound session for a room. *)

    val session_id : t -> string
    (** Get session ID. *)

    val room_id : t -> string
    (** Get room ID. *)

    val message_index : t -> int
    (** Get current message index. *)

    val needs_rotation : t -> bool
    (** Check if session should be rotated. *)

    val export_session_key : t -> string
    (** Export the session key for sharing via m.room_key. *)

    val signing_key : t -> string
    (** Get the signing key. *)

    val encrypt : t -> string -> int * string
    (** Encrypt a message. Returns (message_index, ciphertext). *)

    val mark_shared_with : t -> user_id:string -> device_id:string -> unit
    (** Mark session as shared with a user/device. *)

    val is_shared_with : t -> user_id:string -> device_id:string -> bool
    (** Check if already shared with a user/device. *)

    val shared_with : t -> (string * string) list
    (** Get list of users this session is shared with. *)
  end
end

(** {1 Olm Machine} *)

(** Olm Machine - high-level state machine for E2EE operations. *)
module Machine : sig
  type t
  (** Machine type. *)

  val create : user_id:string -> device_id:string -> t
  (** Create a new OlmMachine. *)

  val identity_keys : t -> string * string
  (** Get identity keys. *)

  val device_keys_for_upload : t -> string * string * string list * (string * string) list
  (** Get device keys for upload. Returns (user_id, device_id, algorithms, keys). *)

  val generate_one_time_keys : t -> int -> unit
  (** Generate one-time keys if needed. *)

  val one_time_keys_for_upload : t -> (string * string * string) list
  (** Get one-time keys for upload. *)

  val mark_keys_as_published : t -> unit
  (** Mark keys as uploaded. *)

  val receive_device_keys : t -> user_id:string -> devices:(string * Keys.queried_device_keys) list -> unit
  (** Store device keys from key query response. *)

  val get_outbound_group_session : t -> room_id:string -> Megolm.Outbound.t
  (** Get or create outbound Megolm session for a room. *)

  val receive_room_key :
    t ->
    sender_key:string ->
    room_id:string ->
    session_id:string ->
    session_key:string ->
    signing_key:string ->
    unit
  (** Store inbound Megolm session from room key event. *)

  val encrypt_room_message : t -> room_id:string -> content:string -> string
  (** Encrypt a room message using Megolm. *)

  val decrypt_room_message :
    t ->
    room_id:string ->
    sender_key:string ->
    session_id:string ->
    ciphertext:string ->
    message_index:int ->
    (string, string) result
  (** Decrypt a room message. *)

  val get_olm_session : t -> their_identity_key:string -> Session.t option
  (** Get or create Olm session for a device. *)

  val create_olm_session :
    t ->
    their_identity_key:string ->
    their_one_time_key:string ->
    Session.t
  (** Create outbound Olm session. *)

  val create_inbound_session :
    t ->
    their_identity_key:string ->
    their_ephemeral_key:string ->
    one_time_key_id:string ->
    Session.t
  (** Process inbound Olm message to create session. *)

  val encrypt_to_device :
    t ->
    their_identity_key:string ->
    their_one_time_key:string ->
    plaintext:string ->
    int * string
  (** Encrypt a to-device message. *)

  val decrypt_to_device :
    t ->
    their_identity_key:string ->
    message_type:int ->
    ciphertext:string ->
    (string, string) result
  (** Decrypt a to-device message. *)

  val one_time_keys_count : t -> int
  (** Number of one-time keys remaining. *)

  val should_upload_keys : t -> bool
  (** Should upload more one-time keys? *)
end