open Polycell
open FrCore
include Common

type input = Polycell.t Event.t
type output = unit

type 'a event = (input, 'a Event.t) sf
type 'a behavior = (input, 'a) sf

type time = Unique.id
let time_context = Unique.create ()

type context = {
	mutable free_id: int;
	mutable sfs: (int * (input, output) sf) list;
	mutable locked: bool;
	queue: (context -> unit) Queue.t;
}

let lock_fun = ref ignore
let unlock_fun = ref ignore

let sync f a =
	!lock_fun ();
	let r =
		try f a with e ->
			!unlock_fun ();
			raise e in
	!unlock_fun ();
	r

let enqueue c f =
	if not (sync (fun c ->
			let locked = c.locked in
			if locked then Queue.add f c.queue
			else c.locked <- true;
			locked) c)
	then begin
		f c;
		try while true do
			(sync (fun c ->
				try Queue.take c.queue
				with Queue.Empty ->
					c.locked <- false;
					raise Exit) c) c
		done with Exit -> ()
	end

let create_context () =
	{ free_id = 0; sfs = []; queue = Queue.create (); locked = false }

(* should this not be reentrant? *)
let copy_context c = { c with free_id = c.free_id; queue = Queue.create (); locked = false }

let _time_now = ref (Unique.gen time_context)

let cycle context input =
	let time' = !_time_now in
	_time_now := Unique.gen time_context;
	context.sfs <-
		List.fold_left (fun sfs' (id, sf) ->
				let sf', _ = reactimate sf input in
				if is_constant sf then sfs'
				else (id, sf') :: sfs')
			[] context.sfs;
	_time_now := time'

(* TODO: re-figure this stuff out *)
let fix_stack = Stack.create ()
let fix_defer f arg =
	try f arg
	with Undefined_signal -> Stack.push (f, arg) fix_stack
let fix_unwind () =
	try while not (Stack.is_empty fix_stack) do
		let f, arg = Stack.top fix_stack in
		(try f arg with Undefined_signal -> raise Exit);
		ignore (Stack.pop fix_stack)
	done with Exit -> ()


let memoize f =
	(* here we should use a stack that mirrors the reentrancy stack so that
	   memoizations aren't lost *BUT* since events are only memoized when
	   they occur, and behaviors remain the same, our memoizations aren't
	   lost (this is of course assuming any changes in the subcontext don't
	   overlap those of the parent) *)
	let memo = ref None in
	fun i -> match !memo with
	| Some (i', o) when i' =* i -> o
	| _ -> let o = f i in memo := Some (i, o); o

let memoizeq f =
	let memo = ref None in
	fun i -> match !memo with
	| Some (i', o) when i' == i -> o
	| _ -> let o = f i in memo := Some (i, o); o

let memoize_t f =
	let f = memoizeq (fun (i, t) -> f i) in
	fun i -> f (i, !_time_now)


let never = never

let now: unit event = now ()

let tag_e v e =
	e >>> arr (Event.tag v)

let map_e f e =
	e >>> arr (Event.map (memoize_t f))

let filter_e f e =
	e >>> arr (Event.filter (memoize_t f))

let zip_e a b =
	(a &&& b) >>> arr2 Event.join

let unzip_e e =
	e >>> arr (Event.map fst), e >>> arr (Event.map snd)

let par_e l =
	par l >>> arr Event.cat

let merge_e l =
	par l >>> arr (List.fold_left Event.lmerge Event.none)

let snapshot_e s e =
	(s &&& e) >>> arr2 Event.get

let once_e e =
	e >>> once

let changes b =
	b >>> changes

let changesq b =
	b >>> changesq

let delta f b =
	b >>> delta (memoize_t f)

let when_e b =
	b >>> edge ()

let edge_e b =
	b >>> FrCore.delta (fun a b -> if a = None then b else None) >>> arr Event.may

let accum_e i e =
	e >>> arr (memoize_t (Event.map memoize_t)) >>> accum i

let collect_e f i e =
	e >>> collect (memoize_t f) i

let switch_e i e =
	(identity &&& e) >>> rswitch i

let collapse_e b =
	(identity &&& b) >>> rcollapse

let fix_e g f =
	let s = fix g f in
	fix_unwind (); s

let may_e e =
	e >>> arr Event.may


let snapshot_b b e =
	(b &&& e) >>> arr2 Event.tag

let latch b e =
	(b &&& e) >>> latch

let hold i e =
	e >>> hold i

let gate_b b e =
	(b &&& e) >>> gate

let track i b =
	b >>> track i

let zip_b a b =
	a &&& b

let unzip_b b =
	b >>> arr fst, b >>> arr snd

let par_b = par

let accum_b i e =
	hold i (accum_e i e)

let collect_b f i e =
	hold i (collect_e f i e)

let switch_b = switch_e
let sswitch_b g f i e =
	(identity &&& e) >>> srswitch g f i
	
let collapse_b = collapse_e
let scollapse_b g i b =
	(identity &&& b) >>> srcollapse g i

let fix_b = fix_e


let lift0 a =
	constant a

let lift1 f a =
	a >>> arr (memoize f)

let lift2 f a b =
	(a &&& b) >>> arr (memoize (fun (a, b) -> f a b))

let lift3 f a b c =
	(a &&& b &&& c) >>> arr (memoize (fun ((a, b), c) -> f a b c))

let lift1q f a =
	a >>> arr (memoizeq f)

let lift2q f a b =
	(a &&& b) >>> arr (memoizeq (fun (a, b) -> f a b))

let lift3q f a b c =
	(a &&& b &&& c) >>> arr (memoizeq (fun ((a, b), c) -> f a b c))


type 'a event_receiver = 'a -> Polycell.t
type 'a cell = 'a event_receiver

let event_receiver () =
	let id = gen_id ()
	and unreachable = ref false in
	(* if our receiver is unreachable, then we can go constant *)
	let finalise_receiver _ = unreachable := true in
	let receiver = box id in
	Gc.finalise finalise_receiver receiver;
	let signal = const_arr (fun i ->
		(* make sure the event doesn't get stuck in the "occuring" state
		 * not sure if this is actually possible, but better to be safe *)
		!unreachable && i = Event.none,
		Event.map_filter (unbox id) i) in
	receiver, signal

let send_event_in context r v =
	enqueue context (fun context -> cycle context (Event.return (r v)))

let new_cell i =
	let r, e = event_receiver () in
	r, hold i e

let set_cell_in = send_event_in


let event_source start stop =
	let id = gen_id ()
	and handle = ref None in
	(* if our signal is unreachable, then we can stop the sender
	   NOTE: can be stronger: unreachable -> not in context *)
	let finalise_signal _ =
		match !handle with Some h -> stop h; handle := None | None -> () in
	let receiver = box id in
	let signal = arr (fun i ->
		(* we're being used; make sure our sender is active *)
		if !handle = None then handle := Some (start receiver);
		Event.map_filter (unbox id) i) in
	Gc.finalise finalise_signal signal;
	signal

let behavior_source force start stop =
	let id = gen_id ()
	and handle = ref None in
	(* if our signal is unreachable, then we can stop the sender
	   NOTE: can be stronger: unreachable -> not in context *)
	let finalise_signal _ =
		match !handle with Some h -> stop h; handle := None | None -> () in
	let receiver = box id in
	let signal = arr (fun i ->
		(* we're being used; make sure our sender is active *)
		if !handle = None then begin
			(* our sender is inactive, so we shouldn't be seeing any input *)
			assert (Event.map_filter (unbox id) i = Event.none);
			handle := Some (start receiver);
			Event.return (force ())
		end
		else Event.map_filter (unbox id) i) in
	Gc.finalise finalise_signal signal;
	(* is there a way we can do this without needing a force () here? *)
	hold (force ()) signal


type simple_endpoint = int

let gen_handler_id context =
	let id = context.free_id in
	context.free_id <- context.free_id + 1;
	id

let register_e_in context sf =
	let id = sync gen_handler_id context in
	enqueue context (fun context ->
		(* NOT constant (), or it will be consted *)
		let sf = sf >>> arr ignore in
		fix_defer
			(fun () ->
				let sf', _ = reactimate sf Event.none in
				if not (is_constant sf) then context.sfs <- (id, sf') :: context.sfs)
			());
	id

let register_b_in = register_e_in

let unregister_in context id = enqueue context (fun context ->
	context.sfs <- List.remove_assoc id context.sfs)

module Mixed = struct
	type 'a t = {
		behavior: 'a behavior;
		event: 'a event;
	}

	let construct b e = { behavior = b; event = e }

	let of_event e = {
		behavior = lift0 None;
		event = map_e some e;
	}

	let of_behavior b = {
		behavior = b;
		event = never;
	}

	let with_event m e = { m with event = e }
	let with_behavior m b = { m with behavior = b }

	let behavior_part m = m.behavior
	let event_part m = m.event
	let destruct m = m.behavior, m.event

	let nothing = { behavior = lift0 None; event = never }

	let map f m = {
		behavior = lift1 f m.behavior;
		event = map_e f m.event;
	}

	let zip a b = {
		behavior = zip_b a.behavior b.behavior;
		event = merge_e [
			zip_e a.event b.event;
			zip_e a.event (snapshot_b b.behavior a.event);
			zip_e (snapshot_b a.behavior b.event) b.event ]
	}

	let merge f a b =
		map (uncurry2 f) (zip a b)

	let snapshot m e =
		merge_e [map_e fst (zip_e m.event e); snapshot_b m.behavior e]

	let unzip m =
		let ab, bb = unzip_b m.behavior
		and ae, be = unzip_e m.event in
		{ behavior = ab; event = ae }, { behavior = bb; event = be }

	let par l = {
		behavior = par_b (List.map behavior_part l);
		event = par_e (List.map event_part l);
	}

	let may i m = {
		behavior = track i m.behavior;
		event = may_e m.event;
	}

	let switch i e = {
		behavior = switch_b i.behavior (map_e behavior_part e);
		event = switch_e i.event (map_e event_part e);
	}

	let collapse b = {
		behavior = collapse_b (lift1 behavior_part b);
		event = collapse_e (lift1 event_part b);
	}

	let meta m = {
		behavior = collapse_b (lift1 behavior_part m.behavior);
		event = merge_e [
			may_e (snapshot (switch nothing (map_e (map some) m.event)) m.event);
			collapse_e (lift1 event_part m.behavior) ]
	}

	let fix g f =
		FrCore.fix
			(fun a -> let m = g a in m.behavior &&& m.event)
			(fun s -> f { behavior = s >>> arr fst; event = s >>> arr snd })
end

type 'a mixed = 'a Mixed.t