(*
 * Copyright (C) 2006-2007 XenSource Ltd.
 * Copyright (C) 2008      Citrix Ltd.
 * Author Vincent Hanquez <vincent.hanquez@eu.citrix.com>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation; version 2.1 only. with the special
 * exception on linking described in file LICENSE.
 *
 * 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 Lesser General Public License for more details.
 *)

module Op = struct include Op end
module Packet = struct include Packet end

module BoundedQueue : sig
	type ('a, 'b) t

	(** [create ~capacity ~classify ~limit] creates a queue with maximum [capacity] elements.
	    This is burst capacity, each element is further classified according to [classify],
	    and each class can have its own [limit].
	    [capacity] is enforced as an overall limit.
	    The [limit] can be dynamic, and can be smaller than the number of elements already queued of that class,
	    in which case those elements are considered to use "burst capacity".
	  *)
	val create: capacity:int -> classify:('a -> 'b) -> limit:('b -> int) -> ('a, 'b) t

	(** [clear q] discards all elements from [q] *)
	val clear: ('a, 'b) t -> unit

	(** [can_push q] when [length q < capacity].	*)
	val can_push: ('a, 'b) t -> 'b -> bool

	(** [push e q] adds [e] at the end of queue [q] if [can_push q], or returns [None]. *)
	val push: 'a -> ('a, 'b) t -> unit option

	(** [pop q] removes and returns first element in [q], or raises [Queue.Empty]. *)
	val pop: ('a, 'b) t -> 'a

	(** [peek q] returns the first element in [q], or raises [Queue.Empty].  *)
	val peek : ('a, 'b) t -> 'a

	(** [length q] returns the current number of elements in [q] *)
	val length: ('a, 'b) t -> int

	(** [debug string_of_class q] prints queue usage statistics in an unspecified internal format. *)
	val debug: ('b -> string) -> (_, 'b) t -> string
end = struct
	type ('a, 'b) t =
		{ q: 'a Queue.t
		; capacity: int
		; classify: 'a -> 'b
		; limit: 'b -> int
		; class_count: ('b, int) Hashtbl.t
		}

	let create ~capacity ~classify ~limit =
		{ capacity; q = Queue.create (); classify; limit; class_count = Hashtbl.create 3 }

	let get_count t classification = try Hashtbl.find t.class_count classification with Not_found -> 0

	let can_push_internal t classification class_count =
		Queue.length t.q < t.capacity && class_count < t.limit classification

	let ok = Some ()

	let push e t =
		let classification = t.classify e in
		let class_count = get_count t classification in
		if can_push_internal t classification class_count then begin
			Queue.push e t.q;
			Hashtbl.replace t.class_count classification (class_count + 1);
			ok
		end
		else
			None

	let can_push t classification =
		can_push_internal t classification @@ get_count t classification

	let clear t =
		Queue.clear t.q;
		Hashtbl.reset t.class_count

	let pop t =
		let e = Queue.pop t.q in
		let classification = t.classify e in
		let () = match get_count t classification - 1 with
		| 0 -> Hashtbl.remove t.class_count classification (* reduces memusage *)
		| n -> Hashtbl.replace t.class_count classification n
		in
		e

	let peek t = Queue.peek t.q
	let length t = Queue.length t.q

	let debug string_of_class t =
		let b = Buffer.create 128 in
		Printf.bprintf b "BoundedQueue capacity: %d, used: {" t.capacity;
		Hashtbl.iter (fun packet_class count ->
			Printf.bprintf b "	%s: %d" (string_of_class packet_class) count
		) t.class_count;
		Printf.bprintf b "}";
		Buffer.contents b
end


exception End_of_file
exception Eagain
exception Noent
exception Invalid
exception Reconnect

let _ =
  Callback.register_exception "Xb.Reconnect" Reconnect

type backend_mmap =
{
	mmap: Xenmmap.mmap_interface;     (* mmaped interface = xs_ring *)
	eventchn_notify: unit -> unit; (* function to notify through eventchn *)
	mutable work_again: bool;
}

type backend_fd =
{
	fd: Unix.file_descr;
}

type backend = Fd of backend_fd | Xenmmap of backend_mmap

type partial_buf = HaveHdr of Partial.pkt | NoHdr of int * bytes

(*
	separate capacity reservation for replies and watch events:
	this allows a domain to keep working even when under a constant flood of
	watch events
*)
type capacity = { maxoutstanding: int; maxwatchevents: int }

module Queue = BoundedQueue

type packet_class =
	| CommandReply
	| Watchevent

let string_of_packet_class = function
	| CommandReply -> "command_reply"
	| Watchevent -> "watch_event"

type t =
{
	backend: backend;
	pkt_out: (Packet.t, packet_class) Queue.t;
	mutable partial_in: partial_buf;
	mutable partial_out: string;
	capacity: capacity
}

let to_read con =
	match con.partial_in with
		| HaveHdr partial_pkt -> Partial.to_complete partial_pkt
		| NoHdr   (i, _)    -> i

let debug t =
	Printf.sprintf "XenBus state: partial_in: %d needed, partial_out: %d bytes, pkt_out: %d packets, %s"
		(to_read t)
		(String.length t.partial_out)
		(Queue.length t.pkt_out)
		(BoundedQueue.debug string_of_packet_class t.pkt_out)

let init_partial_in () = NoHdr
	(Partial.header_size (), Bytes.make (Partial.header_size()) '\000')

let reconnect t = match t.backend with
	| Fd _ ->
		(* should never happen, so close the connection *)
		raise End_of_file
	| Xenmmap backend ->
		Xs_ring.close backend.mmap;
		backend.eventchn_notify ();
		(* Clear our old connection state *)
		Queue.clear t.pkt_out;
		t.partial_in <- init_partial_in ();
		t.partial_out <- ""

let queue con pkt = Queue.push pkt con.pkt_out

let read_fd back _con b len =
	let rd = Unix.read back.fd b 0 len in
	if rd = 0 then
		raise End_of_file;
	rd

let read_mmap back _con b len =
	let s = Bytes.make len '\000' in
	let rd = Xs_ring.read back.mmap s len in
	Bytes.blit s 0 b 0 rd;
	back.work_again <- (rd > 0);
	if rd > 0 then
		back.eventchn_notify ();
	rd

let read con b len =
	match con.backend with
	| Fd backfd     -> read_fd backfd con b len
	| Xenmmap backmmap -> read_mmap backmmap con b len

let write_fd back _con b len =
	Unix.write_substring back.fd b 0 len

let write_mmap back _con s len =
	let ws = Xs_ring.write_substring back.mmap s len in
	if ws > 0 then
		back.eventchn_notify ();
	ws

let write con s len =
	match con.backend with
	| Fd backfd     -> write_fd backfd con s len
	| Xenmmap backmmap -> write_mmap backmmap con s len

(* NB: can throw Reconnect *)
let output con =
	(* get the output string from a string_of(packet) or partial_out *)
	let s = if String.length con.partial_out > 0 then
			con.partial_out
		else if Queue.length con.pkt_out > 0 then
			let pkt = Queue.pop con.pkt_out in
			Packet.to_string pkt
		else
			"" in
	(* send data from s, and save the unsent data to partial_out *)
	if s <> "" then (
		let len = String.length s in
		let sz = write con s len in
		let left = String.sub s sz (len - sz) in
		con.partial_out <- left
	);
	(* after sending one packet, partial is empty *)
	con.partial_out = ""

(* we can only process an input packet if we're guaranteed to have room
   to store the response packet *)
let can_input con = Queue.can_push con.pkt_out CommandReply

(* NB: can throw Reconnect *)
let input con =
	if not (can_input con) then None
	else
	let to_read = to_read con in

	(* try to get more data from input stream *)
	let b = Bytes.make to_read '\000' in
	let sz = if to_read > 0 then read con b to_read else 0 in

	(
	match con.partial_in with
	| HaveHdr partial_pkt ->
		(* we complete the data *)
		if sz > 0 then
			Partial.append partial_pkt (Bytes.to_string b) sz;
		if Partial.to_complete partial_pkt = 0 then (
			let pkt = Packet.of_partialpkt partial_pkt in
			con.partial_in <- init_partial_in ();
			Some pkt
		) else None
	| NoHdr (i, buf)      ->
		(* we complete the partial header *)
		if sz > 0 then
			Bytes.blit b 0 buf (Partial.header_size () - i) sz;
		con.partial_in <- if sz = i then
			HaveHdr (Partial.of_string (Bytes.to_string buf)) else NoHdr (i - sz, buf);
		None
	)

let classify t =
	match t.Packet.ty with
	| Op.Watchevent -> Watchevent
	| _ -> CommandReply

let newcon ~capacity backend =
	let limit = function
		| CommandReply -> capacity.maxoutstanding
		| Watchevent -> capacity.maxwatchevents
	in
	{
	backend = backend;
	pkt_out = Queue.create ~capacity:(capacity.maxoutstanding + capacity.maxwatchevents) ~classify ~limit;
	partial_in = init_partial_in ();
	partial_out = "";
	capacity = capacity;
	}

let open_fd fd = newcon (Fd { fd = fd; })

let open_mmap mmap notifyfct =
	(* Advertise XENSTORE_SERVER_FEATURE_RECONNECTION *)
	Xs_ring.set_server_features mmap (Xs_ring.Server_features.singleton Xs_ring.Server_feature.Reconnection);
	newcon (Xenmmap {
		mmap = mmap;
		eventchn_notify = notifyfct;
		work_again = false; })

let close con =
	match con.backend with
	| Fd backend   -> Unix.close backend.fd
	| Xenmmap backend -> Xenmmap.unmap backend.mmap

let is_fd con =
	match con.backend with
	| Fd _   -> true
	| Xenmmap _ -> false

let is_mmap con = not (is_fd con)

let output_len con = Queue.length con.pkt_out
let has_new_output con = Queue.length con.pkt_out > 0
let has_old_output con = String.length con.partial_out > 0

let has_output con = has_new_output con || has_old_output con

let peek_output con = Queue.peek con.pkt_out

let has_partial_input con = match con.partial_in with
	| HaveHdr _ -> true
	| NoHdr (n, _) -> n < Partial.header_size ()
let has_more_input con =
	match con.backend with
	| Fd _         -> false
	| Xenmmap backend -> backend.work_again

let is_selectable con =
	match con.backend with
	| Fd _   -> true
	| Xenmmap _ -> false

let get_fd con =
	match con.backend with
	| Fd backend -> backend.fd
	| Xenmmap _     -> raise (Failure "get_fd")
