type credential_id = string

type json_decoding_error = [  `Json_decoding of string * string * string ]

type decoding_error = [
  json_decoding_error
  | `Base64_decoding of string * string * string
  | `CBOR_decoding of string * string * string
  | `Unexpected_CBOR of string * string * CBOR.Simple.t
  | `Binary_decoding of string * string * Cstruct.t
  | `Attestation_object_decoding of string * string * string
]

type error = [
  decoding_error
  | `Unsupported_key_type of int
  | `Unsupported_algorithm of int
  | `Unsupported_elliptic_curve of int
  | `Unsupported_attestation_format of string
  | `Invalid_public_key of string
  | `Client_data_type_mismatch of string
  | `Origin_mismatch of string * string
  | `Rpid_hash_mismatch of string * string
  | `Missing_credential_data
  | `Signature_verification of string
]

let pp_error ppf = function
  | `Json_decoding (ctx, msg, json) ->
    Fmt.pf ppf "json decoding error in %s: %s (json: %s)" ctx msg json
  | `Base64_decoding (ctx, msg, data) ->
    Fmt.pf ppf "base64 decoding error in %s: %s (data: %s)" ctx msg data
  | `CBOR_decoding (ctx, msg, data) ->
    Fmt.pf ppf "cbor decoding error in %s: %s (data: %s)" ctx msg data
  | `Unexpected_CBOR (ctx, msg, data) ->
    Fmt.pf ppf "unexpected cbor in %s: %s (data: %s)" ctx msg (CBOR.Simple.to_diagnostic data)
  | `Binary_decoding (ctx, msg, data) ->
    Fmt.pf ppf "binary decoding error in %s: %s (data: %a)" ctx msg Cstruct.hexdump_pp data
  | `Attestation_object_decoding (ctx, msg, data) ->
    Fmt.pf ppf "attestation object decoding error in %s: %s (data: %s)" ctx msg data
  | `Unsupported_key_type i ->
    Fmt.pf ppf "unsupported cose key type %d" i
  | `Unsupported_algorithm i ->
    Fmt.pf ppf "unsupported cose algorithm %d" i
  | `Unsupported_elliptic_curve i ->
    Fmt.pf ppf "unsupported cose elliptic curve %d" i
  | `Unsupported_attestation_format fmt ->
    Fmt.pf ppf "unsupported attestation format %s" fmt
  | `Invalid_public_key msg ->
    Fmt.pf ppf "invalid public key %s" msg
  | `Client_data_type_mismatch is ->
    Fmt.pf ppf "client data type mismatch: received %s" is
  | `Origin_mismatch (should, is) ->
    Fmt.pf ppf "origin mismatch: expected %s, received %s" should is
  | `Rpid_hash_mismatch (should, is) ->
    Fmt.pf ppf "rpid hash mismatch: expected %s received %s"
      (Base64.encode_string should) (Base64.encode_string is)
  | `Missing_credential_data -> Fmt.string ppf "missing credential data"
  | `Signature_verification msg -> Fmt.pf ppf "signature verification failed %s" msg

type t = {
  origin : string;
}

type challenge = string

let generate_challenge ?(size = 32) () =
  if size < 16 then invalid_arg "size must be at least 16 bytes";
  let ch = Mirage_crypto_rng.generate size |> Cstruct.to_string in
  ch, Base64.encode_string ch

let challenge_to_string c = c
let challenge_of_string s = Some s

let challenge_equal = String.equal

let b64_dec thing s =
  Result.map_error
    (function `Msg m -> `Base64_decoding (thing, m, s))
    Base64.(decode ~pad:false ~alphabet:uri_safe_alphabet s)

let guard p e = if p then Ok () else Error e

let (>>=) v f = match v with Ok v -> f v | Error _ as e -> e

type base64url_string = string
let base64url_string_of_yojson = function
  | `String b64 ->
    Base64.(decode ~pad:false ~alphabet:uri_safe_alphabet b64)
    |> Result.map_error (function `Msg m -> m)
  | _ -> Error "base64url_string"

let extract_k_i ctx map k =
  Option.to_result ~none:(`Unexpected_CBOR (ctx, "integer key not present: " ^ string_of_int k, `Map map))
    (Option.map snd
      (List.find_opt (fun (l, _) -> match l with `Int i -> i = k | _ -> false) map))

let extract_k_str ctx map k =
  Option.to_result ~none:(`Unexpected_CBOR (ctx, "string key not present: " ^ k, `Map map))
    (Option.map snd
      (List.find_opt (fun (l, _) -> match l with `Text s -> s = k | _ -> false) map))

let extract_int ctx = function
  | `Int i -> Ok i
  | c -> Error (`Unexpected_CBOR (ctx, "not an integer", c))

let extract_bytes ctx = function
  | `Bytes b -> Ok b
  | c -> Error (`Unexpected_CBOR (ctx, "not bytes", c))

let extract_map ctx = function
  | `Map b -> Ok b
  | c -> Error (`Unexpected_CBOR (ctx, "not a map", c))

let extract_array ctx = function
  | `Array b -> Ok b
  | c -> Error (`Unexpected_CBOR (ctx, "not an array", c))

let extract_text ctx = function
  | `Text s -> Ok s
  | c -> Error (`Unexpected_CBOR (ctx, "not a text", c))

let cose_pubkey cbor_data =
   extract_map "cose pubkey" cbor_data >>= fun kv ->
   extract_k_i "cose pubkey kty" kv 1 >>= extract_int "cose pubkey kty" >>= fun kty ->
   guard (kty = 2) (`Unsupported_key_type kty) >>= fun () ->
   extract_k_i "cose pubkey alg" kv 3 >>= extract_int "cose pubkey alg" >>= fun alg ->
   guard (alg = -7) (`Unsupported_algorithm alg) >>= fun () ->
   extract_k_i "cose pubkey crv" kv (-1) >>= extract_int "cose pubkey crv" >>= fun crv ->
   guard (crv = 1) (`Unsupported_elliptic_curve crv) >>= fun () ->
   extract_k_i "cose pubkey x" kv (-2) >>= extract_bytes "cose pubkey x" >>= fun x ->
   extract_k_i "cose pubkey y" kv (-3) >>= extract_bytes "cose pubkey y" >>= fun y ->
   let four = Cstruct.create 1 in Cstruct.set_uint8 four 0 4;
   let cs = Cstruct.concat [ four ; Cstruct.of_string x ; Cstruct.of_string y ] in
   Result.map_error
     (fun e -> `Invalid_public_key (Fmt.to_to_string Mirage_crypto_ec.pp_error e))
     (Mirage_crypto_ec.P256.Dsa.pub_of_cstruct cs)

let decode_partial_cbor ctx data =
  try Ok (CBOR.Simple.decode_partial data)
  with CBOR.Error m -> Error (`CBOR_decoding (ctx, "failed to decode CBOR " ^ m, data))

let decode_cbor ctx data =
  try Ok (CBOR.Simple.decode data)
  with CBOR.Error m -> Error (`CBOR_decoding (ctx, "failed to decode CBOR " ^ m, data))

let guard_length ctx len cs =
  guard (Cstruct.length cs >= len)
    (`Binary_decoding (ctx, "too short (< " ^ string_of_int len ^ ")", cs))

let parse_attested_credential_data data =
  guard_length "attested credential data" 18 data >>= fun () ->
  let aaguid = Cstruct.sub data 0 16 in
  let cid_len = Cstruct.BE.get_uint16 data 16 in
  let rest = Cstruct.shift data 18 in
  guard_length "attested credential data" cid_len rest >>= fun () ->
  let cid, pubkey = Cstruct.split rest cid_len in
  decode_partial_cbor "public key" (Cstruct.to_string pubkey) >>= fun (pubkey, rest) ->
  cose_pubkey pubkey >>= fun pubkey ->
  Ok ((aaguid, cid, pubkey), Cstruct.of_string rest)

let string_keys ctx kv =
  List.fold_right (fun (k, v) acc ->
    match acc, k with
    | Error _ as e, _ -> e
    | Ok xs, `Text t -> Ok ((t, v) :: xs)
    | _, _ -> Error (`Unexpected_CBOR (ctx, "Map does contain non-text keys", `Map kv)))
    kv (Ok [])

let parse_extension_data data =
   decode_partial_cbor "extension data" (Cstruct.to_string data) >>= fun (data, rest) ->
   extract_map "extension data" data >>= fun kv ->
   string_keys "extension data" kv >>= fun kv ->
   Ok (kv, Cstruct.of_string rest)

type auth_data = {
  rpid_hash : Cstruct.t ;
  user_present : bool ;
  user_verified : bool ;
  sign_count : Int32.t ;
  attested_credential_data : (Cstruct.t * Cstruct.t * Mirage_crypto_ec.P256.Dsa.pub) option ;
  extension_data : (string * CBOR.Simple.t) list option ;
}

let flags byte =
  let b i = byte land (1 lsl i) <> 0 in
  b 0, b 2, b 6, b 7

let parse_auth_data data =
  let data = Cstruct.of_string data in
  guard_length "authenticator data" 37 data >>= fun () ->
  let rpid_hash = Cstruct.sub data 0 32 in
  let user_present, user_verified, attested_data_included, extension_data_included =
    flags (Cstruct.get_uint8 data 32)
  in
  let sign_count = Cstruct.BE.get_uint32 data 33 in
  let rest = Cstruct.shift data 37 in
  (if attested_data_included then
     Result.map (fun (d, r) -> Some d, r) (parse_attested_credential_data rest)
   else Ok (None, rest)) >>= fun (attested_credential_data, rest) ->
  (if extension_data_included then
     Result.map (fun (d, r) -> Some d, r) (parse_extension_data rest)
   else Ok (None, rest)) >>= fun (extension_data, rest) ->
  guard (Cstruct.length rest = 0) (`Binary_decoding ("authenticator data", "leftover", rest)) >>= fun () ->
  Ok { rpid_hash ; user_present ; user_verified ; sign_count ; attested_credential_data ; extension_data }

let parse_attestation_statement fmt data =
  match fmt with
  | "none" when data = [] -> Ok None
  | "none" -> Error (`Unexpected_CBOR ("attestion statement", "format is none, map must be empty", `Map data))
  | "fido-u2f" ->
    extract_k_str "attestation statement" data "x5c" >>= extract_array "attestation statement x5c" >>= fun cert ->
    extract_k_str "attestation statement" data "sig" >>= extract_bytes "attestation statemnt sig" >>= fun signature ->
    begin match cert with
      | [ c ] ->
        extract_bytes "attestation statement x5c" c >>= fun c ->
        Result.map_error
          (function `Msg m -> `Attestation_object_decoding ("attestation statement x5c", m, String.escaped c))
          (X509.Certificate.decode_der (Cstruct.of_string c))
      | cs -> Error (`Attestation_object_decoding ("attestation statement x5c", "expected single certificate", String.concat "," (List.map CBOR.Simple.to_diagnostic cs)))
    end >>= fun cert ->
    Ok (Some (cert, signature))
  | x -> Error (`Unsupported_attestation_format x)

let parse_attestation_object data =
  decode_cbor "attestation object" data >>= extract_map "attestation object" >>= fun kv ->
  extract_k_str "attestation object" kv "fmt" >>= extract_text "attestation object fmt" >>= fun fmt ->
  extract_k_str "attestation object" kv "authData" >>= extract_bytes "attestation object authData" >>= fun auth_data ->
  extract_k_str "attestation object" kv "attStmt" >>= extract_map "attestation object attStmt" >>= fun attestation_statement ->
  parse_auth_data auth_data >>= fun auth_data ->
  parse_attestation_statement fmt attestation_statement >>= fun attestation_statement ->
  Ok (auth_data, attestation_statement)

let of_json_or_err thing p json =
  Result.map_error
    (fun msg -> `Json_decoding (thing, msg, Yojson.Safe.to_string json))
    (p json)

let of_json thing p s =
  (try Ok (Yojson.Safe.from_string s)
   with Yojson.Json_error msg ->
     Error (`Json_decoding (thing, msg, s))) >>=
  of_json_or_err thing p

let json_get member = function
  | `Assoc kv as json ->
    List.assoc_opt member kv
    |> Option.to_result ~none:(`Json_decoding (member, "missing key", Yojson.Safe.to_string json))
  | json -> Error (`Json_decoding (member, "non-object", Yojson.Safe.to_string json))

let json_string thing : Yojson.Safe.t -> (string, _) result = function
  | `String s -> Ok s
  | json -> Error (`Json_decoding (thing, "non-string", Yojson.Safe.to_string json))

let json_assoc thing : Yojson.Safe.t -> ((string * Yojson.Safe.t) list, _) result = function
  | `Assoc s -> Ok s
  | json -> Error (`Json_decoding (thing, "non-assoc", Yojson.Safe.to_string json))

(* XXX: verify [origin] is in fact an origin *)
let create origin = { origin }

let rpid t =
  match String.split_on_char '/' t.origin with
  | [ _protocol ; "" ; host ] -> host
  | _ -> assert false

type credential_data = {
  aaguid : string ;
  credential_id : credential_id ;
  public_key : Mirage_crypto_ec.P256.Dsa.pub ;
}

type registration = {
  user_present : bool ;
  user_verified : bool ;
  sign_count : Int32.t ;
  attested_credential_data : credential_data ;
  authenticator_extensions : (string * CBOR.Simple.t) list option ;
  client_extensions : (string * Yojson.Safe.t) list ;
  certificate : X509.Certificate.t option ;
}

type register_response = {
  attestation_object : base64url_string [@key "attestationObject"];
  client_data_json : base64url_string [@key "clientDataJSON"];
} [@@deriving of_yojson]

let register_response_of_string =
  of_json "register response" register_response_of_yojson

let register t response =
  (* XXX: credential.getClientExtensionResults() *)
  let client_data_hash = Mirage_crypto.Hash.SHA256.digest
    (Cstruct.of_string response.client_data_json) in
  begin try Ok (Yojson.Safe.from_string response.client_data_json)
    with Yojson.Json_error msg ->
      Error (`Json_decoding ("clientDataJSON", msg, response.client_data_json))
  end >>= fun client_data ->
  json_get "type" client_data >>= json_string "type" >>=
  (function
    | "webauthn.create" -> Ok ()
    | wrong_typ -> Error (`Client_data_type_mismatch wrong_typ)) >>= fun () ->
  json_get "challenge" client_data >>= json_string "challenge" >>= fun challenge ->
  b64_dec "response.ClientDataJSON.challenge" challenge >>= fun challenge ->
  json_get "origin" client_data >>= json_string "origin" >>= fun origin ->
  guard (String.equal t.origin origin)
    (`Origin_mismatch (t.origin, origin)) >>= fun () ->
  json_get "clientExtensions" client_data >>= json_assoc "clientExtensions" >>= fun client_extensions ->
  parse_attestation_object response.attestation_object >>= fun (auth_data, attestation_statement) ->
  let rpid_hash = Mirage_crypto.Hash.SHA256.digest (Cstruct.of_string (rpid t)) in
  guard (Cstruct.equal auth_data.rpid_hash rpid_hash)
    (`Rpid_hash_mismatch (Cstruct.to_string rpid_hash, Cstruct.to_string auth_data.rpid_hash)) >>= fun () ->
  (* verify user present, user verified flags in auth_data.flags *)
  Option.to_result ~none:`Missing_credential_data
    auth_data.attested_credential_data >>= fun (aaguid, credential_id, public_key) ->
  begin match attestation_statement with
   | None -> Ok None
   | Some (cert, signature) ->
     let pub_cs = Mirage_crypto_ec.P256.Dsa.pub_to_cstruct public_key in
     let sigdata = Cstruct.concat [
       Cstruct.create 1 ; rpid_hash ; client_data_hash ; credential_id ; pub_cs
     ] in
     let pk = X509.Certificate.public_key cert
     and signature = Cstruct.of_string signature
     in
     Result.map_error (function `Msg m -> `Signature_verification m)
       (X509.Public_key.verify `SHA256 ~signature pk (`Message sigdata)) >>= fun () ->
     Ok (Some cert)
  end >>= fun certificate ->
  (* check attestation cert, maybe *)
  (* check auth_data.attested_credential_data.credential_id is not registered ? *)
  let registration =
    let attested_credential_data = {
      aaguid = Cstruct.to_string aaguid ;
      credential_id = Cstruct.to_string credential_id ;
      public_key
    } in
    {
      user_present = auth_data.user_present ;
      user_verified = auth_data.user_verified ;
      sign_count = auth_data.sign_count ;
      attested_credential_data ;
      authenticator_extensions = auth_data.extension_data ;
      client_extensions ;
      certificate ;
    }
  in
  Ok (challenge, registration)

type authentication = {
  user_present : bool ;
  user_verified : bool ;
  sign_count : Int32.t ;
  authenticator_extensions : (string * CBOR.Simple.t) list option ;
  client_extensions : (string * Yojson.Safe.t) list ;
}

type authenticate_response = {
  authenticator_data : base64url_string [@key "authenticatorData"];
  client_data_json : base64url_string [@key "clientDataJSON"];
  signature : base64url_string ;
  userHandle : base64url_string option ;
} [@@deriving of_yojson]

let authenticate_response_of_string =
  of_json "authenticate response" authenticate_response_of_yojson

let authenticate t public_key response =
  let client_data_hash = Mirage_crypto.Hash.SHA256.digest
    (Cstruct.of_string response.client_data_json) in
  begin try Ok (Yojson.Safe.from_string response.client_data_json)
    with Yojson.Json_error msg ->
      Error (`Json_decoding ("clientDataJSON", msg, response.client_data_json))
  end >>= fun client_data ->
  json_get "type" client_data >>= json_string "type" >>=
  (function
    | "webauthn.get" -> Ok ()
    | wrong_typ -> Error (`Client_data_type_mismatch wrong_typ)) >>= fun () ->
  json_get "challenge" client_data >>= json_string "challenge" >>= fun challenge ->
  b64_dec "response.ClientDataJSON.challenge" challenge >>= fun challenge ->
  json_get "origin" client_data >>= json_string "origin" >>= fun origin ->
  guard (String.equal t.origin origin)
    (`Origin_mismatch (t.origin, origin)) >>= fun () ->
  json_get "clientExtensions" client_data >>= json_assoc "clientExtensions" >>= fun client_extensions ->
  parse_auth_data response.authenticator_data >>= fun auth_data ->
  let rpid_hash = Mirage_crypto.Hash.SHA256.digest (Cstruct.of_string (rpid t)) in
  guard (Cstruct.equal auth_data.rpid_hash rpid_hash)
    (`Rpid_hash_mismatch (Cstruct.to_string rpid_hash, Cstruct.to_string auth_data.rpid_hash)) >>= fun () ->
  let sigdata = Cstruct.concat [ Cstruct.of_string response.authenticator_data ; client_data_hash ]
  and signature = Cstruct.of_string response.signature
  in
  Result.map_error (function `Msg m -> `Signature_verification m)
    (X509.Public_key.verify `SHA256 ~signature (`P256 public_key) (`Message sigdata)) >>= fun () ->
  let authentication = {
    user_present = auth_data.user_present ;
    user_verified = auth_data.user_verified ;
    sign_count = auth_data.sign_count ;
    authenticator_extensions = auth_data.extension_data ;
    client_extensions ;
  } in
  Ok (challenge, authentication)