open Wonka_types;
open Wonka_helpers;

type switchMapStateT('a) = {
  mutable outerTalkback: (. talkbackT) => unit,
  mutable innerTalkback: (. talkbackT) => unit,
  mutable innerActive: bool,
  mutable closed: bool,
  mutable ended: bool,
};

let switchMap = f =>
  curry(source =>
    curry(sink => {
      let state: switchMapStateT('a) = {
        outerTalkback: talkbackPlaceholder,
        innerTalkback: talkbackPlaceholder,
        innerActive: false,
        closed: false,
        ended: false,
      };

      let applyInnerSource = innerSource =>
        innerSource((. signal) =>
          switch (signal) {
          | End =>
            state.innerActive = false;
            state.innerTalkback = talkbackPlaceholder;
            if (state.ended) {
              sink(. End);
            };
          | Start(tb) =>
            state.innerActive = true;
            state.innerTalkback = tb;
            tb(. Pull);
          | Push(x) when !state.closed =>
            sink(. Push(x));
            state.innerTalkback(. Pull);
          | Push(_) => ()
          }
        );

      source((. signal) =>
        switch (signal) {
        | End when !state.ended =>
          state.ended = true;
          if (!state.innerActive) {
            sink(. End);
          };
        | End => ()
        | Start(tb) =>
          state.outerTalkback = tb;
          tb(. Pull);
        | Push(x) when !state.ended =>
          if (state.innerActive) {
            state.innerTalkback(. Close);
            state.innerTalkback = talkbackPlaceholder;
          };
          applyInnerSource(f(. x));
          state.outerTalkback(. Pull);
        | Push(_) => ()
        }
      );

      sink(.
        Start(
          (. signal) =>
            switch (signal) {
            | Pull => state.innerTalkback(. Pull)
            | Close =>
              state.innerTalkback(. Close);
              if (!state.ended) {
                state.ended = true;
                state.closed = true;
                state.outerTalkback(. Close);
                state.innerTalkback = talkbackPlaceholder;
              };
            },
        ),
      );
    })
  );

let switchAll = source => switchMap((. x) => x, source);