(** Server-side key backup and recovery.

    This module implements Matrix server-side backup of room keys using the
    m.megolm_backup.v1.curve25519-aes-sha2 backup algorithm.

    Key backup allows clients to store encrypted room keys on the server so
    they can be recovered on new devices or after data loss.

    Note: Due to various known flaws in this algorithm, it is provided mainly
    for backwards compatibility with existing backups. *)

open Mirage_crypto_ec

(** {1 Backup Key Types} *)

(** Signature verification state *)
type signature_state =
  | Missing           (** No signature found *)
  | Invalid           (** Signature is invalid *)
  | Valid_but_not_trusted  (** Valid but signer not trusted *)
  | Valid_and_trusted      (** Valid and signer is trusted *)

let signature_state_trusted = function
  | Valid_and_trusted -> true
  | _ -> false

(** Signature verification result *)
type signature_verification = {
  device_signature : signature_state;
  user_identity_signature : signature_state;
  other_signatures : (string * signature_state) list;  (** device_id -> state *)
}

let signature_verification_trusted v =
  signature_state_trusted v.device_signature ||
  signature_state_trusted v.user_identity_signature ||
  List.exists (fun (_, s) -> signature_state_trusted s) v.other_signatures

let empty_signature_verification = {
  device_signature = Missing;
  user_identity_signature = Missing;
  other_signatures = [];
}

(** Auth data for m.megolm_backup.v1.curve25519-aes-sha2 *)
type megolm_v1_auth_data = {
  public_key : string;  (** Base64 Curve25519 public key *)
  signatures : (string * (string * string) list) list;  (** user_id -> (key_id, sig) *)
}

(** Convert signatures to JSON *)
let signatures_to_json signatures =
  let user_sigs = List.map (fun (user_id, key_sigs) ->
    let inner = List.map (fun (key_id, sig_) ->
      Printf.sprintf {|"%s":"%s"|} key_id sig_
    ) key_sigs |> String.concat "," in
    Printf.sprintf {|"%s":{%s}|} user_id inner
  ) signatures |> String.concat "," in
  "{" ^ user_sigs ^ "}"

(** Encode auth data to JSON string *)
let megolm_v1_auth_data_to_json auth_data =
  Printf.sprintf {|{"public_key":"%s","signatures":%s}|}
    auth_data.public_key
    (signatures_to_json auth_data.signatures)

(** Room key backup info - describes the backup algorithm and parameters *)
type backup_info =
  | Megolm_v1_curve25519_aes_sha2 of megolm_v1_auth_data
  | Other of { algorithm : string; auth_data : Jsont.json }

(** Backup version info from server *)
type backup_version_info = {
  version : string;
  algorithm : string;
  auth_data : Jsont.json;
  count : int;
  etag : string;
}

(** {1 Backup Encryption Key} *)

(** Private key for decrypting backed up room keys *)
type backup_decryption_key = {
  private_key : string;  (** Base64-encoded X25519 private key *)
  public_key : string;   (** Base64-encoded X25519 public key *)
}

(** Public key for encrypting room keys for backup *)
type backup_encryption_key = {
  public_key : string;   (** Base64-encoded X25519 public key *)
  mutable backup_version : string option;
  mutable signatures : (string * (string * string) list) list;
}

(** Generate a new backup key pair *)
let generate_backup_key () =
  let priv, pub = X25519.gen_key () in
  (* X25519.secret is a string, and gen_key returns (secret, public) where public is a string *)
  let priv_bytes = X25519.secret_to_octets priv in
  {
    private_key = Base64.encode_string priv_bytes;
    public_key = Base64.encode_string pub;
  }

(** Create encryption key from decryption key *)
let encryption_key_of_decryption_key (decryption_key : backup_decryption_key) = {
  public_key = decryption_key.public_key;
  backup_version = None;
  signatures = [];
}

(** Create encryption key from base64 public key *)
let encryption_key_of_base64 public_key =
  match Base64.decode public_key with
  | Error _ -> Error "Invalid base64 encoding"
  | Ok bytes ->
      if String.length bytes <> 32 then
        Error "Invalid key length"
      else
        Ok {
          public_key;
          backup_version = None;
          signatures = [];
        }

(** Get the backup algorithm name *)
let backup_algorithm = "m.megolm_backup.v1.curve25519-aes-sha2"

(** {1 Room Key Encryption for Backup} *)

(** Encrypted session data (matches Matrix spec) *)
type encrypted_session_data = {
  ephemeral : string;   (** Base64 ephemeral public key *)
  ciphertext : string;  (** Base64 ciphertext *)
  mac : string;         (** Base64 MAC *)
}

(** Key backup data for a single session *)
type key_backup_data = {
  first_message_index : int;
  forwarded_count : int;
  is_verified : bool;
  session_data : encrypted_session_data;
}

(** Room key backup request (for upload) *)
type keys_backup_request = {
  rooms : (string * (string * key_backup_data) list) list;  (** room_id -> session_id -> data *)
}

(** Encrypt a room key for backup using X25519/AES-256/HMAC-SHA256 *)
let encrypt_room_key (encryption_key : backup_encryption_key) ~session_key ~session_id ~room_id ~sender_key =
  match Base64.decode encryption_key.public_key with
  | Error _ -> Error "Invalid encryption key"
  | Ok recipient_pub_bytes ->
      if String.length recipient_pub_bytes <> 32 then
        Error "Invalid X25519 public key length"
      else
        (* Generate ephemeral key pair *)
        let ephemeral_priv, ephemeral_pub = X25519.gen_key () in

        (* Perform X25519 key exchange *)
        (match X25519.key_exchange ephemeral_priv recipient_pub_bytes with
           | Error _ -> Error "Key exchange failed"
           | Ok shared_secret ->
               (* Derive encryption and MAC keys using HKDF *)
               let prk = Hkdf.extract ~hash:`SHA256 ~salt:"" shared_secret in
               let okm = Hkdf.expand ~hash:`SHA256 ~prk ~info:"" 64 in
               let aes_key = String.sub okm 0 32 in
               let mac_key = String.sub okm 32 32 in

               (* Build the payload to encrypt *)
               let payload = Printf.sprintf
                 {|{"algorithm":"m.megolm.v1.aes-sha2","room_id":"%s","sender_key":"%s","session_id":"%s","session_key":"%s"}|}
                 room_id sender_key session_id session_key
               in

               (* Pad payload to 16-byte boundary (PKCS#7) *)
               let pad_len = 16 - (String.length payload mod 16) in
               let padded = payload ^ String.make pad_len (Char.chr pad_len) in

               (* Generate random IV *)
               let iv = Mirage_crypto_rng.generate 16 in

               (* Encrypt with AES-256-CBC *)
               let key = Mirage_crypto.AES.CBC.of_secret aes_key in
               let ciphertext = Mirage_crypto.AES.CBC.encrypt ~key ~iv padded in

               (* Prepend IV to ciphertext for MAC calculation *)
               let mac_input = iv ^ ciphertext in

               (* Calculate HMAC-SHA256 *)
               let mac = Digestif.SHA256.hmac_string ~key:mac_key mac_input in
               let mac_bytes = Digestif.SHA256.to_raw_string mac in
               (* Truncate to first 8 bytes as per spec *)
               let mac_truncated = String.sub mac_bytes 0 8 in

               Ok {
                 ephemeral = Base64.encode_string ephemeral_pub;
                 ciphertext = Base64.encode_string (iv ^ ciphertext);
                 mac = Base64.encode_string mac_truncated;
               })

(** Decrypt a room key from backup *)
let decrypt_room_key (decryption_key : backup_decryption_key) (session_data : encrypted_session_data) =
  match Base64.decode decryption_key.private_key,
        Base64.decode session_data.ephemeral,
        Base64.decode session_data.ciphertext,
        Base64.decode session_data.mac with
  | Error _, _, _, _ -> Error "Invalid private key encoding"
  | _, Error _, _, _ -> Error "Invalid ephemeral key encoding"
  | _, _, Error _, _ -> Error "Invalid ciphertext encoding"
  | _, _, _, Error _ -> Error "Invalid MAC encoding"
  | Ok priv_bytes, Ok ephemeral_bytes, Ok ciphertext_with_iv, Ok mac_bytes ->
      match X25519.secret_of_octets priv_bytes with
      | Error _ -> Error "Invalid X25519 private key"
      | Ok (priv, _) ->
          (* Perform X25519 key exchange *)
          (match X25519.key_exchange priv ephemeral_bytes with
           | Error _ -> Error "Key exchange failed"
           | Ok shared_secret ->
               (* Derive encryption and MAC keys using HKDF *)
               let prk = Hkdf.extract ~hash:`SHA256 ~salt:"" shared_secret in
               let okm = Hkdf.expand ~hash:`SHA256 ~prk ~info:"" 64 in
               let aes_key = String.sub okm 0 32 in
               let mac_key = String.sub okm 32 32 in

               (* Verify MAC *)
               let expected_mac = Digestif.SHA256.hmac_string ~key:mac_key ciphertext_with_iv in
               let expected_mac_bytes = Digestif.SHA256.to_raw_string expected_mac in
               let expected_mac_truncated = String.sub expected_mac_bytes 0 8 in
               if mac_bytes <> expected_mac_truncated then
                 Error "MAC verification failed"
               else begin
                 (* Extract IV and ciphertext *)
                 if String.length ciphertext_with_iv < 16 then
                   Error "Ciphertext too short"
                 else begin
                   let iv = String.sub ciphertext_with_iv 0 16 in
                   let ciphertext = String.sub ciphertext_with_iv 16 (String.length ciphertext_with_iv - 16) in

                   (* Decrypt with AES-256-CBC *)
                   let key = Mirage_crypto.AES.CBC.of_secret aes_key in
                   let plaintext = Mirage_crypto.AES.CBC.decrypt ~key ~iv ciphertext in

                   (* Remove PKCS#7 padding *)
                   let pad_len = Char.code (String.get plaintext (String.length plaintext - 1)) in
                   if pad_len < 1 || pad_len > 16 then
                     Error "Invalid padding"
                   else
                     let unpadded = String.sub plaintext 0 (String.length plaintext - pad_len) in
                     Ok unpadded
                 end
               end)

(** {1 Backup Machine State} *)

(** State of the backup machine *)
type backup_state =
  | Disabled           (** No backup configured *)
  | Creating           (** Creating a new backup *)
  | Enabling           (** Enabling existing backup *)
  | Resuming           (** Resuming existing backup *)
  | Enabled            (** Backup is active *)
  | Downloading        (** Downloading keys from backup *)
  | Disabling          (** Disabling backup *)

(** Backup machine for managing room key backups *)
type t = {
  user_id : Matrix_proto.Id.User_id.t;
  device_id : Matrix_proto.Id.Device_id.t;
  mutable state : backup_state;
  mutable encryption_key : backup_encryption_key option;
  mutable decryption_key : backup_decryption_key option;
  mutable backup_version : string option;
  (* Pending sessions to backup *)
  mutable pending_sessions : (string * string * string) list;  (** (room_id, session_id, sender_key) *)
  (* Sessions that have been backed up *)
  mutable backed_up_sessions : (string * string) list;  (** (room_id, session_id) *)
}

(** Create a new backup machine *)
let create ~user_id ~device_id = {
  user_id;
  device_id;
  state = Disabled;
  encryption_key = None;
  decryption_key = None;
  backup_version = None;
  pending_sessions = [];
  backed_up_sessions = [];
}

(** Check if backup is enabled *)
let is_enabled t = t.state = Enabled

(** Get the current backup version *)
let backup_version t = t.backup_version

(** {1 Backup Setup} *)

(** Enable backup with a new key *)
let enable_with_new_key t =
  let key = generate_backup_key () in
  t.decryption_key <- Some key;
  t.encryption_key <- Some (encryption_key_of_decryption_key key);
  t.state <- Creating;
  key

(** Enable backup with an existing decryption key *)
let enable_with_key t decryption_key =
  t.decryption_key <- Some decryption_key;
  t.encryption_key <- Some (encryption_key_of_decryption_key decryption_key);
  t.state <- Enabling

(** Enable backup with only an encryption key (upload-only mode) *)
let enable_upload_only t encryption_key version =
  t.encryption_key <- Some encryption_key;
  t.decryption_key <- None;
  t.backup_version <- Some version;
  encryption_key.backup_version <- Some version;
  t.state <- Enabled

(** Set the backup version after creating *)
let set_backup_version t version =
  t.backup_version <- Some version;
  (match t.encryption_key with
   | Some key -> key.backup_version <- Some version
   | None -> ());
  t.state <- Enabled

(** Disable backup *)
let disable t =
  t.state <- Disabling;
  t.encryption_key <- None;
  t.decryption_key <- None;
  t.backup_version <- None;
  t.backed_up_sessions <- [];
  t.state <- Disabled

(** {1 Session Management} *)

(** Mark a session as needing backup *)
let mark_session_for_backup t ~room_id ~session_id ~sender_key =
  if not (List.mem (room_id, session_id) t.backed_up_sessions) then
    t.pending_sessions <- (room_id, session_id, sender_key) :: t.pending_sessions

(** Get number of pending sessions *)
let pending_count t = List.length t.pending_sessions

(** Check if a session has been backed up *)
let is_session_backed_up t ~room_id ~session_id =
  List.mem (room_id, session_id) t.backed_up_sessions

(** Mark a session as backed up *)
let mark_session_backed_up t ~room_id ~session_id =
  t.pending_sessions <- List.filter (fun (r, s, _) -> r <> room_id || s <> session_id) t.pending_sessions;
  if not (List.mem (room_id, session_id) t.backed_up_sessions) then
    t.backed_up_sessions <- (room_id, session_id) :: t.backed_up_sessions

(** {1 Room Key Recovery} *)

(** Recovered room key data *)
type recovered_room_key = {
  room_id : string;
  session_id : string;
  session_key : string;
  sender_key : string;
  algorithm : string;
  forwarded : bool;
}

(** Parse a recovered room key from decrypted JSON *)
let parse_recovered_key json_str =
  (* Simple JSON parsing - in production would use proper parser *)
  let _get_field name str =
    let _pattern = Printf.sprintf {|"%s":"|} name in
    match String.split_on_char '"' str with
    | _ -> None
  in
  (* Simplified - would use Jsont in real implementation *)
  match Jsont_bytesrw.decode_string Jsont.json json_str with
  | Error _ -> None
  | Ok _json -> None  (* Would extract fields from JSON *)

(** Result of importing keys from backup *)
type import_result = {
  imported_count : int;
  total_count : int;
  keys : recovered_room_key list;
}

(** {1 Backup Creation API Helpers} *)

(** Build auth data for a new backup *)
let build_auth_data encryption_key = {
  public_key = encryption_key.public_key;
  signatures = encryption_key.signatures;
}

(** Create backup version request body *)
let create_version_request_body t =
  match t.encryption_key with
  | None -> Error "No encryption key configured"
  | Some key ->
      let auth_data = build_auth_data key in
      let auth_json = megolm_v1_auth_data_to_json auth_data in
      Ok (Printf.sprintf
        {|{"algorithm":"%s","auth_data":%s}|}
        backup_algorithm auth_json)

(** {1 Recovery Key Format} *)

(** Encode a backup decryption key as a recovery key (human-readable) *)
let encode_recovery_key (key : backup_decryption_key) =
  match Base64.decode key.private_key with
  | Error _ -> Error "Invalid private key"
  | Ok bytes ->
      (* Add header byte 0x8B, then calculate parity byte *)
      let with_header = "\x8B" ^ bytes in
      let parity = String.fold_left (fun acc c -> acc lxor Char.code c) 0 with_header in
      let full = with_header ^ String.make 1 (Char.chr parity) in
      (* Encode as base58 with spaces every 4 chars for readability *)
      let _base58_alphabet = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz" in
      (* Simplified - would use proper base58 encoding; using base64 for now *)
      let encoded = Base64.encode_string ~pad:false full in
      (* Add spaces for readability *)
      let with_spaces = String.to_seq encoded
        |> Seq.mapi (fun i c -> if i > 0 && i mod 4 = 0 then [' '; c] else [c])
        |> Seq.flat_map List.to_seq
        |> String.of_seq
      in
      Ok with_spaces

(** Decode a recovery key to a backup decryption key *)
let decode_recovery_key recovery_key =
  (* Remove spaces and decode *)
  let cleaned = String.split_on_char ' ' recovery_key |> String.concat "" in
  match Base64.decode cleaned with
  | Error _ -> Error "Invalid recovery key format"
  | Ok bytes ->
      if String.length bytes < 3 then
        Error "Recovery key too short"
      else if String.get bytes 0 <> '\x8B' then
        Error "Invalid recovery key header"
      else begin
        (* Verify parity *)
        let key_bytes = String.sub bytes 1 (String.length bytes - 2) in
        let expected_parity = Char.code (String.get bytes (String.length bytes - 1)) in
        let actual_parity = String.fold_left (fun acc c -> acc lxor Char.code c) 0
          (String.sub bytes 0 (String.length bytes - 1)) in
        if expected_parity <> actual_parity then
          Error "Recovery key parity check failed"
        else begin
          (* Derive public key from private key *)
          match X25519.secret_of_octets key_bytes with
          | Error _ -> Error "Invalid private key in recovery key"
          | Ok (priv, pub_bytes) ->
              ignore priv;
              Ok {
                private_key = Base64.encode_string key_bytes;
                public_key = Base64.encode_string pub_bytes;
              }
        end
      end