(*********************************************************************************)
(*                OCaml-IRI                                                      *)
(*                                                                               *)
(*    Copyright (C) 2016 Institut National de Recherche en Informatique          *)
(*    et en Automatique. All rights reserved.                                    *)
(*                                                                               *)
(*    This program is free software; you can redistribute it and/or modify       *)
(*    it under the terms of the GNU Lesser General Public License version        *)
(*    3 as published by the Free Software Foundation.                            *)
(*                                                                               *)
(*    This program is distributed in the hope that it will be useful,            *)
(*    but WITHOUT ANY WARRANTY; without even the implied warranty of             *)
(*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the              *)
(*    GNU Library General Public License for more details.                       *)
(*                                                                               *)
(*    You should have received a copy of the GNU Lesser General Public           *)
(*    License along with this program; if not, write to the Free Software        *)
(*    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA                   *)
(*    02111-1307  USA                                                            *)
(*                                                                               *)
(*    Contact: Maxence.Guesdon@inria.fr                                          *)
(*                                                                               *)
(*                                                                               *)
(*********************************************************************************)

type path =
| Absolute of string list
| Relative of string list

module KV = Map.Make(String)
type query_kv = string KV.t

type t = {
    scheme : string ;
    user : string option ;
    host : string option ;
    port : int option ;
    path : path ;
    query : string option ; (** not %-decoded as query is not parsed to name/value pairs *)
    fragment : string option ;
    query_kv : query_kv  option ;
      (** key value pairs from query string to avoid parsing it various times *)
  }

type error =
| Parse_error of string * exn

exception Error of error

let is_absolute t =
  match t.fragment with
    None -> t.scheme <> ""
  | _ -> false

let is_relative t = t.scheme = ""

let utf8_nb_bytes_of_char c =
  let n = Char.code c in
  if n < 0b10000000 then
    1
  else if n < 0b11100000 then
      2
    else if n < 0b11110000 then
        3
      else
        4
let is_ucschar c =
  let n = Uchar.to_int c in
  (n >= 0xA0 && n <= 0xD7FF) ||
    (n >= 0xF900 && n<= 0xFDCF) ||
    (n >= 0xFDF0 && n <= 0xFFEF) ||
    (n >= 0x10000 && n<= 0x1FFFD) ||
    (n >= 0x20000 && n<= 0x2FFFD) ||
    (n >= 0x30000 && n<= 0x3FFFD) ||
    (n >= 0x40000 && n<= 0x4FFFD) ||
    (n >= 0x50000 && n<= 0x5FFFD) ||
    (n >= 0x60000 && n<= 0x6FFFD) ||
    (n >= 0x70000 && n<= 0x7FFFD) ||
    (n >= 0x80000 && n<= 0x8FFFD) ||
    (n >= 0x90000 && n<= 0x9FFFD) ||
    (n >= 0xA0000 && n<= 0xAFFFD) ||
    (n >= 0xB0000 && n<= 0xBFFFD) ||
    (n >= 0xC0000 && n<= 0xCFFFD) ||
    (n >= 0xD0000 && n<= 0xDFFFD) ||
    (n >= 0xE1000 && n<= 0xEFFFD)
let is_iprivate c =
  let n = Uchar.to_int c in
  (n >= 0xE000 && n <= 0xF8FF) ||
    (n >= 0xF0000 && n <= 0xFFFFD) ||
    (n >= 0x100000 && n <= 0x10FFFD)

let pct_decode =
  let rec iter b len s i =
    if i >= len then
      ()
    else
      begin
        let i =
          match s.[i] with
            '%' when i+2 < len ->
              begin
                try
                  let n = int_of_string ("0x"^(String.sub s (i+1) 2)) in
                  let c = Char.chr n in
                  Buffer.add_char b c;
                  i+3
                with
                  _ -> Buffer.add_char b s.[i]; i+1
              end
          | _ ->
              let size = utf8_nb_bytes_of_char s.[i] in
              if size = 1 then
                Buffer.add_char b s.[i]
              else
                Buffer.add_substring b s i size ;
              i+size
        in
        iter b len s i
      end
  in
  fun s ->
    let len = String.length s in
    let b = Buffer.create len in
    iter b len s 0;
    Buffer.contents b
;;

let safe_chars =
  let f n =
    match Char.chr n with
      'A'..'Z' | 'a'..'z' | '0'..'9' | '-' | '.' -> true
    | _ -> false
  in
  Array.init 256 f
;;

let from_safe_chars ?(f=fun _ -> false) safe_chars c =
  let n = Uchar.to_int c in
  assert (n >= 0);
  if n > 255 then f c else safe_chars.(n)

let scheme_safe_chars =
  let a = Array.copy safe_chars in
  a.(Char.code '+') <- true;
  a
;;
let scheme_safe_char = from_safe_chars scheme_safe_chars

let sub_delims =
  [| '!' ; '$' ; '&' ; '\'' ; '(' ; ')' ;
     '*' ; '+' ; ',' ; ';' ; '=' |]
;;

let user_safe_chars =
  let a = Array.copy safe_chars in
  a.(Char.code ':') <- true;
  a.(Char.code '~') <- true;
  a.(Char.code '_') <- true;
  Array.iter (fun c -> a.(Char.code c) <- true) sub_delims ;
  a
;;
let user_safe_char =
  from_safe_chars  ~f: is_ucschar user_safe_chars

let host_safe_chars =
  let a = Array.copy safe_chars in
  a.(Char.code '~') <- true;
  a.(Char.code '_') <- true;
  a.(Char.code '[') <- true;
  a.(Char.code ']') <- true;
  a.(Char.code ':') <- true;
  Array.iter (fun c -> a.(Char.code c) <- true) sub_delims ;
  a
;;
let host_safe_char =
  from_safe_chars ~f: is_ucschar host_safe_chars

let path_safe_chars =
  let a = Array.copy safe_chars in
  a.(Char.code '~') <- true;
  a.(Char.code '_') <- true;
  a.(Char.code ':') <- true;
  a.(Char.code '@') <- true;
  Array.iter (fun c -> a.(Char.code c) <- true) sub_delims ;
  a
;;
let path_safe_char =
  from_safe_chars ~f: is_ucschar path_safe_chars

let query_part_safe_chars =
  let a = Array.copy safe_chars in
  a.(Char.code '~') <- true;
  a.(Char.code '_') <- true;
  a.(Char.code ':') <- true;
  a.(Char.code '@') <- true;
  a.(Char.code '?') <- true;
  a.(Char.code '/') <- true;
  Array.iter (fun c -> a.(Char.code c) <- true) sub_delims ;
  a.(Char.code '=') <- false;
  a.(Char.code '&') <- false;
  a
;;
let query_part_safe_char =
  from_safe_chars ~f: is_ucschar query_part_safe_chars

let fragment_safe_chars =
  let a = Array.copy safe_chars in
  a.(Char.code '~') <- true;
  a.(Char.code '_') <- true;
  a.(Char.code ':') <- true;
  a.(Char.code '@') <- true;
  a.(Char.code '?') <- true;
  a.(Char.code '/') <- true;
  Array.iter (fun c -> a.(Char.code c) <- true) sub_delims ;
  a
;;
let fragment_safe_char =
  from_safe_chars ~f: is_ucschar fragment_safe_chars

let pct_encode_utf8 dest codepoint =
  let b = Buffer.create 4 in
  Uutf.Buffer.add_utf_8 b codepoint;
  let s = Buffer.contents b in
  String.iter
    (fun c -> Printf.bprintf dest "%%%02X" (Char.code c))
    s

let pct_encode_b =
  let f is_safe_char b () _i = function
    `Malformed str -> Buffer.add_string b str
  | `Uchar codepoint ->
      if is_safe_char codepoint then
        Uutf.Buffer.add_utf_8 b codepoint
      else
        pct_encode_utf8 b codepoint
  in
  fun b is_safe_char s ->
    Uutf.String.fold_utf_8 (f is_safe_char b) () s
;;

let pct_encode is_safe_char s =
  let b = Buffer.create (String.length s) in
  pct_encode_b b is_safe_char s ;
  Buffer.contents b

let pct_encode_query =
  let safe_chars =
    let a = Array.copy safe_chars in
    a.(Char.code '~') <- true;
    a.(Char.code '_') <- true;
    a.(Char.code ':') <- true;
    a.(Char.code '@') <- true;
    a.(Char.code '?') <- true;
    a.(Char.code '/') <- true;
    Array.iter (fun c -> a.(Char.code c) <- true) sub_delims ;
    a
  in
  let is_safe_char =
    from_safe_chars ~f: is_ucschar safe_chars
  in
  fun str -> pct_encode is_safe_char str

let path_string =
  let string_of_path encode l =
    let l =
      if encode then
        List.map (pct_encode path_safe_char) l
      else
        l
    in
    String.concat "/" l
  in
  fun ?(pctencode=false) t ->
    match t.path with
      Absolute l -> "/"^(string_of_path pctencode l)
    | Relative l -> string_of_path pctencode l

let to_string =
  fun ?(pctencode=true) iri  ->
    let has_ihier = iri.host <> None in
    let b = Buffer.create 256 in
    if iri.scheme <> "" then
      (
       Buffer.add_string b iri.scheme ;
       Buffer.add_string b ":"
      );
    if has_ihier then Buffer.add_string b "//";
    (match iri.user with
       None -> ()
     | Some u ->
         Buffer.add_string b
           (if pctencode then pct_encode user_safe_char u else u);
         Buffer.add_char b '@'
    );
    (match iri.host with
       None -> ()
     | Some s ->
         Buffer.add_string b
           (if pctencode then pct_encode host_safe_char s else s)
    );
    (match iri.port with None -> () | Some n -> Buffer.add_string b (":"^(string_of_int n))) ;

    Buffer.add_string b (path_string ~pctencode iri);

    (match iri.query with
       None -> ()
     | Some s ->
         Buffer.add_char b '?';
         Buffer.add_string b s ;
    );
    (match iri.fragment with
       None -> ()
     | Some s ->
         Buffer.add_char b '#' ;
         Buffer.add_string b (if pctencode then pct_encode fragment_safe_char s else s)
    );
    Buffer.contents b
;;

let pp fmt iri = Format.pp_print_string fmt (to_string iri)

let map_opt f = function None -> None | Some x -> Some (f x)

let utf8_split =
  let f split_char b_chars acc_words _i = function
    `Malformed str -> Buffer.add_string b_chars str; acc_words
  | `Uchar codepoint ->
      if split_char codepoint then
        begin
          let w = Buffer.contents b_chars in
          Buffer.reset b_chars ;
          match w with
            "" -> acc_words
          | _ -> w :: acc_words
        end
      else
        begin
          Uutf.Buffer.add_utf_8 b_chars codepoint ;
          acc_words
        end
  in
  fun split_char str ->
    let b = Buffer.create 128 in
    let words = Uutf.String.fold_utf_8 (f split_char b) [] str in
    match Buffer.contents b with
      "" -> List.rev words
    | w -> List.rev (w :: words)

let encode_query_string_part = pct_encode query_part_safe_char

let split_query_string =
  let dec = pct_decode in
  let cp_equal = Uchar.of_char '=' in
  let cp_amp = Uchar.of_char '&' in
  let is_equal = Uchar.equal cp_equal in
  let is_amp = Uchar.equal cp_amp in
  let add map str =
    match utf8_split is_equal str with
      [] | [_] -> KV.add (dec str) "" map
    | [ k ; v ] -> KV.add (dec k) (dec v) map
    | k :: vals -> KV.add (dec k) (dec (String.concat "=" vals)) map
  in
  fun str ->
    let l = utf8_split is_amp str in
    List.fold_left add KV.empty l

let split_query_opt = function
| None -> KV.empty
| Some q -> split_query_string q

let query_string_of_kv =
  let f b max_i i (k, v) =
     pct_encode_b b query_part_safe_char k ;
     Buffer.add_char b '=' ;
     pct_encode_b b query_part_safe_char v ;
     if i <> max_i then Buffer.add_char b '&'
  in
  fun ?q pairs ->
    let b = Buffer.create 128 in
    (match q with
       None | Some "" -> ()
     | Some str ->
         Buffer.add_string b str ;
         Buffer.add_char b '&'
    );
    let pairs = KV.bindings pairs in
    List.iteri (f b (List.length pairs - 1)) pairs ;
    Buffer.contents b

let query_string_of_kv_opt = map_opt query_string_of_kv

let iri ?(scheme="") ?user ?host ?port ?(path=Absolute[]) ?query_kv ?query ?fragment () =
  let (query, query_kv) =
    match query, query_kv with
    | None, None -> (None, None)
    | Some qs, None -> (query, Some (split_query_string qs))
    | None, Some kv -> (Some (query_string_of_kv kv), query_kv)
    | Some q, Some kv ->
        (* concat both and set kv pairs to none so that it will
           be recreated from the concatenation *)
        let qs = query_string_of_kv ~q kv in
        (Some qs, Some (split_query_string qs))
  in
  { scheme ; user ; host ; port ; path ; query ; fragment ; query_kv }

let scheme t = t.scheme
let with_scheme t scheme = { t with scheme }

let user t = t.user
let with_user t user = { t with user }

let host t = t.host
let with_host t host = { t with host }

let port t = t.port
let with_port t port = { t with port }

let path t = t.path
let with_path t path = { t with path }

let append_path t strings =
  let l = match t.path with
    Absolute l | Relative l -> l
  in
  let l =
    match List.rev l with
      "" :: q -> (List.rev q) @ strings
    | _ -> l @ strings
  in
  let path = match t.path with
    | Absolute _-> Absolute l
    | Relative _ -> Relative l
  in
  with_path t path

let query t = t.query
let with_query t query =
  match query with
  | None -> { t with query = None ; query_kv = None }
  | Some qs -> { t with query ; query_kv = Some (split_query_string qs) }

let query_kv t = Option.value ~default:KV.empty t.query_kv

let with_query_kv t map =
  if KV.is_empty map then
    { t with query = None ; query_kv = None }
  else
    { t with
      query = Some (query_string_of_kv map) ;
      query_kv = Some map ;
    }

let query_get t key =
  match KV.find_opt key (query_kv t) with
  | None -> ""
  | Some str -> str

let query_opt t key = KV.find_opt key (query_kv t)

let query_set t k v =
  let map = query_kv t in
  let map = KV.add k v map in
  with_query_kv t map

let fragment t = t.fragment
let with_fragment t fragment = { t with fragment }

let remove_dot_segments =
  let rec iter acc = function
  | [] -> List.rev acc
  | ["."] -> List.rev ("" :: acc)
  | "." :: q -> iter acc q
  | ".." :: q ->
      begin
        match acc with
        | [] -> iter acc q
        | _ :: acc when q = [] -> List.rev ("" :: acc)
        | _ :: acc -> iter acc q
      end
  | h :: q -> iter (h :: acc) q
  in
  iter []

let normalize_path =
  let rec iter acc  = function
  | [] -> List.rev acc
  | [""] -> List.rev (""::acc)
  | "" :: q -> iter acc q
  | h :: q -> iter (h :: acc) q
  in
  fun l -> iter [] (remove_dot_segments l)
;;

let path_normalize = function
| Absolute l -> Absolute (normalize_path l)
| Relative l -> Relative (normalize_path l)

let path_remove_dot_segments = function
  Absolute l -> Absolute (remove_dot_segments l)
| Relative l -> Relative (remove_dot_segments l)

let remove_dot_segments t =
  let path = path_remove_dot_segments t.path in
  { t with path }

let normalize_host s =
  let len = String.length s in
  if len > 0 then
    match String.get s 0 with
      '[' -> String.uppercase_ascii s (* uppercase hexa *)
    | _ -> String.lowercase_ascii s (* lowercase regname *)
  else
    s

let normalize_port t =
  match String.lowercase_ascii t.scheme, t.port with
    "http", Some 80 -> { t with port = None }
  | "https", Some 443 -> { t with port = None }
  | "ftp", Some 21 -> { t with port = None }
  | "sftp", Some 115 -> { t with port = None }
  | "ssh", Some 22 -> { t with port = None }
  | "smtp", Some 25 -> { t with port = None }
  | _ -> t

let normalize_case t =
  { t with
    scheme = String.lowercase_ascii t.scheme ;
    host = map_opt normalize_host t.host ;
  }

let normalize_nfkc t =
  let f = Uunf_string.normalize_utf_8 `NFKC in
  let path =
    match t.path with
      Absolute l -> Absolute (List.map f l)
    | Relative l -> Relative (List.map f l)
  in
  let query = map_opt f t.query (* beware: the query is not %-decoded *) in
  let query_kv = match query with
    | None -> None
    | Some str -> Some (split_query_string str)
  in
  { t with
    host = map_opt f t.host ;
    path ;
    user = map_opt f t.user ;
    query ;
    fragment = map_opt f t.fragment ;
    query_kv ;
  }

let normalize ?(nfkc=true) t =
  let t = { t with path = path_normalize t.path } in
  let t = normalize_case t in
  let t = normalize_port t in
  if nfkc then normalize_nfkc t else t

let compare =
  let compare_scheme i1 i2 = String.compare i1.scheme i2.scheme in
  let compare_str_opt s1 s2 =
    match s1, s2 with
    | None, None -> 0
    | None, _ -> -1
    | Some _, None -> 1
    | Some s1, Some s2 -> String.compare s1 s2
  in
  let compare_host i1 i2 = compare_str_opt i1.host i2.host in
  let compare_user i1 i2 = compare_str_opt i1.user i2.user in
  let compare_port i1 i2 =
    match i1.port, i2.port with
    | None, None -> 0
    | None, _ -> -1
    | Some _, None -> 1
    | Some p1, Some p2 -> Int.compare p1 p2
  in
  let compare_path i1 i2 =
    match i1.path, i2.path with
    | Relative p1, Relative p2 -> List.compare String.compare p1 p2
    | Relative _, Absolute _ -> -1
    | Absolute _, Relative _ -> 1
    | Absolute p1, Absolute p2 ->
        match p1, p2 with
        | [], [""]
        | [""], [] -> 0
        | _ -> List.compare String.compare p1 p2
  in
  let compare_query i1 i2 =
    match i1.query_kv, i2.query_kv with
    | Some q1, Some q2 -> KV.compare String.compare q1 q2
    | None, None -> 0
    | None, _ -> -1
    | _, None -> 1
  in
  let compare_fragment i1 i2 = compare_str_opt i1.fragment i2.fragment in
  let compare_funs = [
      compare_path ;
      compare_host ;
      compare_query ;
      compare_fragment ;
      compare_scheme ;
      compare_user ;
      compare_port ;
    ]
  in
  let compare =
    let rec iter i1 i2 = function
    | [] -> 0
    | f :: q ->
        match f i1 i2 with
        | 0 -> iter i1 i2 q
        | n -> n
    in
    iter
  in
  let norm = normalize in
  fun ?(normalize=false) i1 i2 ->
    let (i1, i2) =
      if normalize
      then norm i1, norm i2
      else i1, i2
    in
    compare i1 i2 compare_funs

(*let compare i1 =
  ignore(query_kv i1);
  let i1 = normalize i1 in
  fun i2 ->
    ignore(query_kv i2);
    let i2 = normalize i2 in
    Pervasives.compare i1 i2
*)
let equal ?normalize i1 i2 = compare ?normalize i1 i2 = 0

let to_uri =
  let user_safe_char = from_safe_chars user_safe_chars in
  let host_safe_char = from_safe_chars host_safe_chars in
  let fragment_safe_char = from_safe_chars fragment_safe_chars in
  let query_part_safe_char = from_safe_chars query_part_safe_chars in
  let path_safe_char = from_safe_chars path_safe_chars in
  let path_string =
    let string_of_path l =
      let l = List.map (pct_encode path_safe_char) l in
      String.concat "/" l
    in
    fun t ->
      match t.path with
        Absolute l -> "/"^(string_of_path l)
      | Relative l -> string_of_path l
  in
  let print_query_string =
    let f b max_i i (k, v) =
      pct_encode_b b query_part_safe_char k ;
      Buffer.add_char b '=' ;
      pct_encode_b b query_part_safe_char v ;
      if i <> max_i then Buffer.add_char b '&'
    in
    fun b t ->
      match KV.bindings (query_kv t) with
      | [] -> ()
      | pairs ->
          Buffer.add_char b '?';
          List.iteri (f b (List.length pairs - 1)) pairs
  in
  fun iri  ->
    let has_ihier = iri.host <> None in
    let b = Buffer.create 256 in
    if iri.scheme <> "" then
      (
       Buffer.add_string b iri.scheme ;
       Buffer.add_string b ":"
      );
    if has_ihier then Buffer.add_string b "//";
    (match iri.user with
       None -> ()
     | Some u ->
         pct_encode_b b user_safe_char u ;
         Buffer.add_char b '@'
    );
    (match iri.host with
       None -> ()
     | Some s -> pct_encode_b b host_safe_char s
    );
    (match iri.port with
     | None -> ()
     | Some n -> Buffer.add_string b (":"^(string_of_int n))
    ) ;

    Buffer.add_string b (path_string iri);

    print_query_string b iri ;

    (match iri.fragment with
       None -> ()
     | Some s ->
         Buffer.add_char b '#' ;
         pct_encode_b b fragment_safe_char s
    );
    Buffer.contents b
;;

let to_string_details =
  let opt_s = function None -> "None" | Some s -> Printf.sprintf "%S" s in
  let opt_n = function None -> "None" | Some n -> Printf.sprintf "%d" n in
  fun iri ->
    let path =
      let s, l = match path iri with
        | Absolute l -> "Absolute", l
        | Relative l -> "Relative", l
      in
      Printf.sprintf "%s [%s]" s (String.concat " ; " (List.map (fun s -> Printf.sprintf "%S" s) l))
    in
    Printf.sprintf "scheme=%S, user=%s, host=%s, port=%s, path=%s, query=%s, fragment=%s"
      (scheme iri) (opt_s (user iri)) (opt_s (host iri))
      (opt_n (port iri)) path (opt_s (query iri)) (opt_s (fragment iri))

let pp_details fmt iri = Format.pp_print_string fmt (to_string_details iri)
