open Lwt.Syntax
open Markup

exception InvalidStanza of string

type element = {
    namespace : string;
    local_name : string;
    attributes : (string * string) list;
    children : (element, string) Either.t list;
  }

(** [tree s] is an [element option] representing the XML element inside of stream [s],
    if [s] is a complete element from start to end. *)
let tree s : element option =
  let element (namespace, name) attributes children =
    Either.Left {
        namespace;
        local_name=name;
        attributes=List.filter_map
                     (fun ((ns, name), content) ->
                       (* remove xmlns -- we don't need it. *)
                       match ns with
                       | "http://www.w3.org/2000/xmlns/" -> None
                       | _ -> Some (name, content))
                     attributes;
        children;
      }
  and text ss = Either.Right (String.concat "" ss) in
  let opt_el = tree ~text ~element s in
  Option.bind opt_el (Either.fold ~left:Option.some ~right:(fun _ -> None))

(** [element_to_string element] is a string representation of the underlying XML in
    [element], for debugging purposes.

    Note this isn't serialization: namely, XML namespaces are inferred and don't exist
    in the actual [element]. *)
let element_to_string ?(indent = 2) (el : element) =
  let rec element_to_string parent {local_name; attributes; children; namespace} =
    let attributes =
      (if parent == namespace then "" else " xmlns=\"" ^ namespace ^ "\"") ^
        String.concat ""
          (List.map (fun (n, a) -> " " ^ n ^ "=\"" ^ a ^ "\"") attributes)
    in
    let tab s =
      String.split_on_char '\n' s |> String.concat ("\n" ^ String.make indent ' ')
    in
    let tag = "\n<" ^ local_name ^ attributes
    in tag ^ match children with
             | [] -> "/>"
             | _ -> ">" ^ (children
                           |> List.map
                                (Either.fold
                                   ~left:(element_to_string namespace)
                                   ~right:(fun x -> "\n" ^ x))
                           |> String.concat ""
                           |> tab)
                    ^ "\n</" ^ local_name ^ ">"
  in element_to_string "" el |> String.trim

(** [next stream] is a promise containing a full stanza of the fragments of
    [stream]. *)
let next (stream : (signal, async) stream) : (signal, sync) stream Lwt.t =
  let traverse_stanza depth fragment =
    let depth = match fragment with
      | `Start_element _ -> depth + 1
      | `End_element -> depth - 1
      | _ -> depth
    in ([fragment], if depth = 0 then None else Some depth)
  in transform traverse_stanza 0 stream |> Markup_lwt.load

(** [get stream] is a promise containing a single Xml element of [stream]. *)
let get (stream : (signal, async) stream) : element Lwt.t =
  let* signal = next stream
  in match tree signal with
     | Some xml -> Lwt.return xml
     | None -> Lwt.fail (InvalidStanza (signal |> write_xml |> to_string))