use entire peak (technically) when rendering info box HTML, clean up

module Genetics.Browser.Track.UI where

import Prelude

import Control.Coroutine (Consumer, Producer, connect, runProcess)
import Control.Coroutine as Co
import Control.Monad.Aff (Aff, Milliseconds, delay, forkAff, killFiber, launchAff, launchAff_)
import Control.Monad.Aff.AVar (AVar)
import Control.Monad.Aff.AVar as AVar
import Control.Monad.Eff (Eff)
import Control.Monad.Eff.Class (liftEff)
import Control.Monad.Eff.Console (log)
import Control.Monad.Eff.Exception (error, throw)
import DOM.Classy.HTMLElement (appendChild, setId) as DOM
import DOM.Classy.ParentNode (toParentNode)
import DOM.HTML (window) as DOM
import DOM.HTML.Types (htmlDocumentToDocument) as DOM
import DOM.HTML.Window (document) as DOM
import DOM.Node.Document (createElement, documentElement) as DOM
import DOM.Node.ParentNode (querySelector) as DOM
import DOM.Node.Types (Element, ElementId)
import Data.Array as Array
import Data.Bifunctor (bimap)
import Data.BigInt (BigInt)
import Data.BigInt as BigInt
import Data.Either (Either(Right, Left))
import Data.Filterable (filterMap)
import Data.Foldable (class Foldable, foldMap, length, null)
import Data.Foreign (Foreign, MultipleErrors, renderForeignError)
import Data.Generic.Rep (class Generic)
import Data.Generic.Rep.Show (genericShow)
import Data.Int as Int
import Data.Lens ((^.))
import Data.Lens as Lens
import Data.Map (Map)
import Data.Map as Map
import Data.Maybe (Maybe(..), fromJust, fromMaybe)
import Data.Monoid (class Monoid, mempty)
import Data.Newtype (over, unwrap, wrap)
import Data.Pair (Pair(..))
import Data.Pair as Pair
import Data.Record.Extra (eqRecord)
import Data.Symbol (SProxy(SProxy))
import Data.Traversable (for_, traverse_)
import Data.Tuple (Tuple(Tuple))
import Data.Variant (Variant, case_, inj)
import Data.Variant as V
import Genetics.Browser.Track.Backend (Peak, RenderedTrack, pixelSegments)
import Genetics.Browser.Track.Demo (Annotation, AnnotationField, SNP, annotationsForScale, demoBrowser, filterSig, getAnnotations, getGenes, getSNPs, showAnnotationField)
import Genetics.Browser.Track.UI.Canvas (BrowserCanvas, TrackPadding, _Dimensions, _Track, browserCanvas, browserOnClick, debugBrowserCanvas, dragScroll, renderBrowser, setBrowserCanvasSize, setElementStyle, wheelZoom)
import Genetics.Browser.Types (ChrId(ChrId), _NegLog10)
import Genetics.Browser.Types.Coordinates (CoordSys, CoordSysView(CoordSysView), _TotalSize, coordSys, normalizeView, pairSize, pairsOverlap, pixelsView, scaleViewBy, translateViewBy, viewScale)
import Global.Unsafe (unsafeStringify)
import Graphics.Canvas as Canvas
import Graphics.Drawing (Drawing, Point)
import Math as Math
import Partial.Unsafe (unsafePartial)
import Simple.JSON (read)
import Unsafe.Coerce (unsafeCoerce)

foreign import windowInnerSize :: ∀ e. Eff e Canvas.Dimensions

-- | Set an event to fire on the given button id
foreign import buttonEvent :: ∀ e.
                              String
                           -> Eff e Unit
                           -> Eff e Unit

-- | Set callback to run after the window has resized (see UI.js for
-- | time waited), providing it with the new window size.
foreign import resizeEvent :: ∀ e.
                              ({ width  :: Number
                               , height :: Number } -> Eff e Unit)
                           -> Eff e Unit

-- | The value returned by the function that initializes the genome browser;
-- | provides an interface to the running browser instance.
type BrowserInterface e a
  = { getView          :: Aff e CoordSysView
    , getBrowserCanvas :: Aff e BrowserCanvas
    , lastOverlaps     :: Aff e (Number -> Point -> a)
    , queueCommand     :: Variant UICmdR -> Aff e Unit
    , queueUpdateView  :: UpdateView -> Eff e Unit
    }

-- | Creates the browser using the provided initial data, returning
-- | a BrowserInterface for reading state & sending commands to it
initializeBrowser :: CoordSys _ _
                  -> (BrowserCanvas -> CoordSysView
                      -> { tracks :: { snps :: RenderedTrack (SNP ())
                                     , annotations :: RenderedTrack (Annotation ()) }
                         , relativeUI :: Drawing
                         , fixedUI :: Drawing })
                  -> CoordSysView
                  -> BrowserCanvas
                  -> Aff _ (BrowserInterface _ ({snps :: Array (SNP ())}))
initializeBrowser cSys mainBrowser initView bc = do

  renderFiberVar <- AVar.makeEmptyVar

  viewVar <- AVar.makeVar initView
  bcVar <- AVar.makeVar bc

  uiCmdVar <- AVar.makeEmptyVar
  lastOverlapsVar <- AVar.makeEmptyVar

  drawCachedBrowser <- browserCache mainBrowser

  let getView = AVar.readVar viewVar
      getBrowserCanvas = AVar.readVar bcVar
      queueCommand = flip AVar.putVar uiCmdVar
      lastOverlaps = AVar.readVar lastOverlapsVar

    -- hardcoded timeout for now
  queueUpdateView <- uiViewUpdate cSys (wrap 100.0) {view: viewVar, uiCmd: uiCmdVar}

  let mainLoop = do
        uiCmd <- AVar.takeVar uiCmdVar
        case_ # V.on _render (\_ -> pure unit)
              # V.on _docResize (\ {width} -> do
                      canvas <- AVar.takeVar bcVar
                      let {height} = canvas ^. _Dimensions
                      canvas' <- liftEff $ setBrowserCanvasSize {width, height} canvas
                      AVar.putVar canvas' bcVar
                      queueCommand $ inj _render unit
                      mainLoop)
              $ uiCmd

        -- if there's a rendering fiber running, we kill it
        traverse_ (killFiber (error "Resetting renderer"))
          =<< AVar.tryTakeVar renderFiberVar

        csView <- getView
        canvas <- getBrowserCanvas
        toRender <- drawCachedBrowser canvas csView

        let currentScale = viewScale (canvas ^. _Track <<< _Dimensions) csView
            (Pair offset _) = pixelsView currentScale csView

                          -- offset the click by the current canvas translation
        _ <- AVar.tryTakeVar lastOverlapsVar
        AVar.putVar (\n r -> let r' = {x: r.x + offset, y: r.y }
                             in { snps: toRender.tracks.snps.overlaps n r' }) lastOverlapsVar

        -- fork a new renderFiber
        renderFiber <- forkAff
                       $ renderBrowser (wrap 2.0) canvas offset toRender
        AVar.putVar renderFiber renderFiberVar

        mainLoop

  _ <- forkAff mainLoop

  queueCommand $ inj _render unit

  pure { getView
       , getBrowserCanvas
       , lastOverlaps
       , queueCommand
       , queueUpdateView }

data UpdateView =
    ScrollView Number
  | ZoomView Number
  | ModView (Pair BigInt -> Pair BigInt)

instance showUpdateView :: Show UpdateView where
  show (ScrollView x) = "(Scroll by " <> show x <> ")"
  show (ZoomView s) = "(Zoom by " <> show s <> ")"
  show _ = "(ModView)"

    -- TODO idk if this instance makes sense??? whatevs
instance semigroupUpdateView :: Semigroup UpdateView where
  append (ScrollView x1) (ScrollView x2) = ScrollView (x1 + x2)
  append (ZoomView s1)   (ZoomView s2)   = ZoomView   (s1 * s2)
  append _ y = y

instance monoidUpdateView :: Monoid UpdateView where
  mempty = ModView id

updateViewFold :: UpdateView
               -> CoordSysView
               -> CoordSysView
updateViewFold uv iv = case uv of
  ZoomView   x -> iv `scaleViewBy`     x
  ScrollView x -> iv `translateViewBy` x
  ModView f    -> over CoordSysView f iv

queueCmd :: ∀ a. AVar a -> a -> Eff _ Unit
queueCmd av cmd = launchAff_ $ AVar.putVar cmd av

-- | Provided with the AVar containing the view state (which is updated atomically)
-- | and the UICmd AVar queue (which is only written to),
-- | return a function that queues view updates using the provided timeout.
uiViewUpdate :: ∀ r.
                CoordSys _ BigInt
             -> Milliseconds
             -> { view  :: AVar CoordSysView
                , uiCmd :: AVar (Variant UICmdR) | r }
             -> Aff _ (UpdateView -> Eff _ Unit)
uiViewUpdate cs timeout { view, uiCmd } = do
  viewCmd <- AVar.makeEmptyVar

  curCmdVar <- AVar.makeVar (ModView id)

  let normView = normalizeView cs (BigInt.fromInt 200000)
      loop' updater = do
        cmd <- AVar.takeVar viewCmd

        killFiber (error "Resetting view update") updater
        liftEff $ log $ "forking view update"

        curCmd <- AVar.takeVar curCmdVar
        let cmd' = curCmd <> cmd
        AVar.putVar cmd' curCmdVar

        updater' <- forkAff do
            delay timeout
            liftEff $ log "Running view update"

            vr <- AVar.takeVar view
            let vr' = normView $ updateViewFold cmd' vr
            AVar.putVar vr' view

            AVar.putVar (inj _render unit) uiCmd
            AVar.takeVar curCmdVar *> AVar.putVar mempty curCmdVar

        loop' updater'

  _ <- forkAff $ loop' (pure unit)

  pure $ queueCmd viewCmd

btnUI :: { scrollMod :: Number, zoomMod :: Number }
      -> (UpdateView -> Eff _ Unit)
      -> Eff _ Unit
btnUI mods cb = do
  buttonEvent "scrollLeft"  $ cb $ ScrollView     (-mods.scrollMod)
  buttonEvent "scrollRight" $ cb $ ScrollView       mods.scrollMod
  buttonEvent "zoomOut"     $ cb $ ZoomView $ 1.0 + mods.zoomMod
  buttonEvent "zoomIn"      $ cb $ ZoomView $ 1.0 - mods.zoomMod

type UICmdR = ( render :: Unit
              , docResize :: { width :: Number, height :: Number } )

_render = SProxy :: SProxy "render"
_docResize = SProxy :: SProxy "docResize"

debugView :: BrowserInterface _ _
          -> Eff _ { get :: _
                   , set :: _ }
debugView s = do
  let get name = launchAff_ do
         view <- s.getView
         liftEff do
           log $ "CoordSysView: " <> show (map BigInt.toString $ unwrap view)
           setWindow name $ (\(Pair l r) -> {l,r}) $ unwrap view

  let set lr = s.queueUpdateView
               $ ModView $ const (BigInt.fromNumber <$> Pair lr.l lr.r)

  pure {get, set}

browserCache
  :: ∀ a b.
     (BrowserCanvas
      -> CoordSysView
      -> { tracks :: { snps :: a
                     , annotations :: b }
         , relativeUI :: Drawing
         , fixedUI :: Drawing })
     -> Aff _ (BrowserCanvas
               -> CoordSysView
               -> Aff _ ({ tracks :: { snps :: a
                                     , annotations :: b }
                         , relativeUI :: Drawing
                         , fixedUI :: Drawing }))
browserCache f = do
  bcDim <- AVar.makeEmptyVar
  viewSize <- AVar.makeEmptyVar

  lastPartial <- AVar.makeEmptyVar
  lastFinal <- AVar.makeEmptyVar

  pure \bc csv -> do

    let newBCDim = bc ^. _Dimensions
    oldBCDim <- AVar.tryTakeVar bcDim
    AVar.putVar newBCDim bcDim

    oldPartial <- AVar.tryTakeVar lastPartial

    partial <- case oldBCDim, oldPartial of
      Just d, Just o -> do
        let changed = not $ d `eqRecord` newBCDim
        liftEff $ log $ "BrowserCanvas changed? " <> (show changed)
        -- remove the cached final output value on canvas resize,
        -- since we need to recalculate everything
        when changed (void $ AVar.tryTakeVar lastFinal)
        pure $ if d `eqRecord` newBCDim then o else f bc

      _, _ -> do
        liftEff $ log "Cache was empty! Recalculating all"
        _ <- AVar.tryTakeVar lastFinal
        pure $ f bc

    AVar.putVar partial lastPartial

    let newViewSize = pairSize $ unwrap csv
    oldViewSize <- AVar.tryTakeVar viewSize
    AVar.putVar newViewSize viewSize

    oldFinal <- AVar.tryTakeVar lastFinal

    let output = case oldViewSize, oldFinal of
          Just vs, Just o -> if newViewSize == vs then o else partial csv
          _, _ -> partial csv

    AVar.putVar output lastFinal
    pure output

printSNPInfo :: ∀ r. Array (SNP r) -> Eff _ Unit
printSNPInfo fs = do
  let n = length fs :: Int
      m = 5
  log $ "showing " <> show m <> " out of " <> show n <> " clicked glyphs"
  for_ (Array.take m fs) (log <<< unsafeCoerce)

wrapWith :: String -> String -> String
wrapWith tag x =
  "<"<> tag <>">" <> x <> "</"<> tag <>">"

snpHTML :: ∀ r.
           SNP r
        -> String
snpHTML {position, feature} = wrapWith "div" contents
  where contents = foldMap (wrapWith "p")
            [ "SNP: "    <> feature.name
            , "Chr: "    <> show feature.chrId
            , "Pos: "    <> show (Pair.fst position)
            , "Score: "  <> show feature.score
            , "-log10: " <> show (unwrap $ (feature.score ^. _NegLog10))
            ]

peakHTML :: ∀ a.
            (a -> String)
         -> Peak _ _ a
         -> String
peakHTML disp peak =
  case Array.uncons peak.elements of
    Nothing               -> ""
    Just {head, tail: []} -> disp head
    Just {head, tail}     ->
      wrapWith "div" $ wrapWith "p"
        $ show (length tail + 1) <> " annotations"

-- | Given a function to transform the data in the annotation's "rest" field
-- | to text (or Nothing if the field should not be displayed), produce a
-- | function that generates HTML from annotations
annotationHTML :: (AnnotationField -> Maybe String)
                -> Annotation () -> String
annotationHTML disp {feature} = wrapWith "div" contents
  where url = fromMaybe "No URL"
              $ map (\a -> "URL: <a target='_blank' href='"
                           <> a <> "'>" <> a <> "") feature.url

        name = fromMaybe ("Annotated SNP: " <> feature.name)
                         (("Gene: " <> _) <$> feature.gene)

        showOther fv = fv.field <> ": " <> (unsafeCoerce fv.value)

        contents = foldMap (wrapWith "p")
          $ [ name
            , "Pos: " <> show (feature.pos)
            , url
            , "Other data: "
            ] <> (filterMap disp
                  $ Array.fromFoldable feature.rest)

-- | Shows all data in "rest" using the default showAnnotationField (which uses unsafeCoerce)
annotationHTMLAll :: Annotation () -> String
annotationHTMLAll =
  annotationHTML (pure <<< showAnnotationField)

-- | Example HTML generator that only shows the "anno" field
annotationHTMLAnnoOnly :: Annotation () -> String
annotationHTMLAnnoOnly = annotationHTML disp
  where disp fv
          | fv.field == "anno" = pure $ showAnnotationField fv
          | otherwise          = Nothing

foreign import initDebugDiv :: ∀ e. Number -> Eff e Unit
foreign import setDebugDivVisibility :: ∀ e. String -> Eff e Unit
foreign import setDebugDivPoint :: ∀ e. Point -> Eff e Unit

foreign import setElementContents :: ∀ e. Element -> String -> Eff e Unit

data InfoBoxF
  = IBoxShow
  | IBoxHide
  | IBoxSetY Int
  | IBoxSetX Int
  | IBoxSetContents String

derive instance genericInfoBoxF :: Generic InfoBoxF _

instance showInfoBoxF :: Show InfoBoxF where
  show = genericShow

updateInfoBox :: Element -> InfoBoxF -> Eff _ Unit
updateInfoBox el cmd =
  case cmd of
    IBoxShow ->
      setElementStyle el "visibility" "visible"
    IBoxHide ->
      setElementStyle el "visibility" "hidden"
    (IBoxSetX x)    ->
      setElementStyle el "left" $ show x <> "px"
    (IBoxSetY y)    ->
      setElementStyle el "top"  $ show y <> "px"
    (IBoxSetContents html) ->
      setElementContents el html

infoBoxId :: ElementId
infoBoxId = wrap "infoBox"

initInfoBox :: Eff _ (InfoBoxF -> Eff _ Unit)
initInfoBox = do
  doc <- map DOM.htmlDocumentToDocument
           $ DOM.document =<< DOM.window
  el <- DOM.createElement "div" doc

  DOM.setId infoBoxId el

  DOM.documentElement doc >>= case _ of
    Nothing -> throw "Couldn't find document body!"
    Just docBody -> void $ DOM.appendChild el docBody

  pure $ updateInfoBox el

runBrowser :: Conf -> BrowserCanvas -> Eff _ _
runBrowser config bc = launchAff $ do

  let clickRadius = 1.0

  liftEff $ initDebugDiv clickRadius

  cmdInfoBox <- liftEff $ initInfoBox

  let cSys :: CoordSys ChrId BigInt
      cSys = coordSys mouseChrSizes

  trackData <-
    {genes:_, snps:_, annotations:_}
    <$> foldMap (getGenes       cSys) config.urls.genes
    <*> foldMap (getSNPs        cSys) config.urls.snps
    <*> foldMap (getAnnotations cSys) config.urls.annotations

  let mainBrowser = demoBrowser cSys config trackData

  browser <-
    initializeBrowser cSys mainBrowser (wrap $ Pair zero (cSys^._TotalSize)) bc

  liftEff do
    resizeEvent \d ->
      launchAff_ $ browser.queueCommand $ inj _docResize d

    btnUI { scrollMod: 0.05, zoomMod: 0.1 } browser.queueUpdateView

    dragScroll bc \ {x,y} ->
       when (Math.abs x >= one) $ launchAff_ do
         trackDims <- Lens.view (_Track <<< _Dimensions)
                      <$> browser.getBrowserCanvas

         liftEff $ browser.queueUpdateView
                 $ ScrollView $ (-x) / trackDims.width

    let scrollZoomScale = 0.06
    wheelZoom bc \dY ->
       browser.queueUpdateView
         $ ZoomView $ 1.0 + scrollZoomScale * dY

  liftEff do
    let overlayDebug :: _
        overlayDebug p = do
          setDebugDivVisibility "visible"
          setDebugDivPoint p

    let sigSnps = filterSig config.score trackData.snps
        annotAround pks snp =
          Array.find (\a -> a.covers `pairsOverlap` snp.position)
            =<< Map.lookup snp.feature.chrId pks

        glyphClick :: _
        glyphClick p = launchAff_ do
          v  <- browser.getView
          bc <- browser.getBrowserCanvas

          let segs = pixelSegments cSys bc v
              annoPeaks = annotationsForScale cSys sigSnps
                            trackData.annotations segs

          lastOverlaps' <- browser.lastOverlaps
          liftEff do
            let clicked = (lastOverlaps' clickRadius p).snps
            printSNPInfo clicked
            case Array.head clicked of
              Nothing -> cmdInfoBox IBoxHide
              Just g  -> do
                cmdInfoBox IBoxShow
                cmdInfoBox $ IBoxSetX $ Int.round p.x
                cmdInfoBox $ IBoxSetContents
                  $ snpHTML g
                  <> foldMap (peakHTML annotationHTMLAll) (annotAround annoPeaks g)

    browserOnClick bc
      { overlay: \_ -> pure unit
      -- { overlay: overlayDebug
      , track:   glyphClick }

  pure unit

type DataURLs = { snps        :: Maybe String
                , annotations :: Maybe String
                , genes       :: Maybe String
                }

type Conf = { browserHeight :: Number
            , trackPadding :: TrackPadding
            , score :: { min :: Number, max :: Number, sig :: Number }
            , urls :: DataURLs
            }

foreign import setWindow :: ∀ e a. String -> a -> Eff e Unit

main :: Foreign -> Eff _ _
main rawConfig = do

  el' <- do
    doc <- DOM.htmlDocumentToDocument
           <$> (DOM.document =<< DOM.window)
    DOM.querySelector (wrap "#browser") (toParentNode doc)

  case read rawConfig :: Either MultipleErrors Conf of
    Left errs -> traverse_ (log <<< renderForeignError) errs
    Right c   -> do
      case el' of
        Nothing -> log "Could not find element '#browser'"
        Just el -> do

          {width} <- windowInnerSize
          let dimensions = { width, height: c.browserHeight }
          bc <- browserCanvas dimensions c.trackPadding el

          debugBrowserCanvas "debugBC" bc

          log $ unsafeStringify c
          void $ runBrowser c bc

-- TODO this could almost certainly be done better
chunkConsumer :: ∀ m a.
                 Foldable m
              => Monoid (m a)
              => AVar (m a)
              -> Consumer (m a) (Aff _) Unit
chunkConsumer av = Co.consumer \m ->
  if null m
    then pure $ Just unit
    else do
      sofar <- AVar.takeVar av
      AVar.putVar (sofar <> m) av
      delay (wrap 20.0)
      pure Nothing

type TrackVar a = AVar (Map ChrId (Array a))
type TrackProducer eff a = Producer (Map ChrId (Array a)) (Aff eff) Unit

-- Feels like this one takes care of a bit too much...
fetchLoop1 :: ∀ a.
              (Maybe (Aff _ (TrackProducer _ a)))
           -> Aff _ (TrackVar a)
fetchLoop1 Nothing = AVar.makeEmptyVar
fetchLoop1 (Just startProd) = do
  prod <- startProd
  avar <- AVar.makeVar mempty
  _ <- forkAff $ runProcess $ prod `connect` chunkConsumer avar
  pure avar

-- | Starts threads that fetch & parse each of the provided tracks,
-- | filling an AVar over time per track, which can be used by other parts of the application
-- | (read only, should be a newtype)
{-
fetchLoop :: CoordSys ChrId BigInt
          -> DataURLs
          -> Aff _
               { gwas :: TrackVar (GWASFeature ())
               , genes :: TrackVar BedFeature
               }
fetchLoop cs urls = do
  gwas <-        fetchLoop1 $ produceGWAS   cs <$> urls.gwas
  genes <-       fetchLoop1 $ produceGenes  cs <$> urls.genes
  pure { gwas, genes }
-}

mouseChrSizes :: Array (Tuple ChrId BigInt)
mouseChrSizes =
  unsafePartial
  $ map (bimap ChrId (fromJust <<< BigInt.fromString))
      [ Tuple "1"   "195471971"
      , Tuple "2"   "182113224"
      , Tuple "3"   "160039680"
      , Tuple "4"   "156508116"
      , Tuple "5"   "151834684"
      , Tuple "6"   "149736546"
      , Tuple "7"   "145441459"
      , Tuple "8"   "129401213"
      , Tuple "9"   "124595110"
      , Tuple "10"  "130694993"
      , Tuple "11"  "122082543"
      , Tuple "12"  "120129022"
      , Tuple "13"  "120421639"
      , Tuple "14"  "124902244"
      , Tuple "15"  "104043685"
      , Tuple "16"  "98207768"
      , Tuple "17"  "94987271"
      , Tuple "18"  "90702639"
      , Tuple "19"  "61431566"
      , Tuple "X"   "171031299"
      , Tuple "Y"   "91744698"
      ]

module Genetics.Browser.Track.UI where

import Prelude

import Color (black)
import Control.Coroutine (Consumer, Producer, connect, runProcess)
import Control.Coroutine as Co
import Control.Monad.Aff (Aff, Milliseconds, delay, forkAff, killFiber, launchAff, launchAff_)
import Control.Monad.Aff.AVar (AVar)
import Control.Monad.Aff.AVar as AVar
import Control.Monad.Eff (Eff)
import Control.Monad.Eff.Class (liftEff)
import Control.Monad.Eff.Console (log)
import Control.Monad.Eff.Exception (error, throw)
import DOM.Classy.HTMLElement (appendChild, setId) as DOM
import DOM.Classy.ParentNode (toParentNode)
import DOM.HTML (window) as DOM
import DOM.HTML.Types (htmlDocumentToDocument) as DOM
import DOM.HTML.Window (document) as DOM
import DOM.Node.Document (createElement, documentElement) as DOM
import DOM.Node.ParentNode (querySelector) as DOM
import DOM.Node.Types (Element, ElementId)
import Data.Array as Array
import Data.Bifunctor (bimap)
import Data.BigInt (BigInt)
import Data.BigInt as BigInt
import Data.Either (Either(Right, Left))
import Data.Filterable (filterMap)
import Data.Foldable (class Foldable, foldMap, length, null)
import Data.Foreign (Foreign, MultipleErrors, renderForeignError)
import Data.Generic.Rep (class Generic)
import Data.Generic.Rep.Show (genericShow)
import Data.Int as Int
import Data.Lens ((^.))
import Data.Lens as Lens
import Data.Map (Map)
import Data.Map as Map
import Data.Maybe (Maybe(..), fromJust, fromMaybe)
import Data.Monoid (class Monoid, mempty)
import Data.Newtype (over, unwrap, wrap)
import Data.Pair (Pair(..))
import Data.Pair as Pair
import Data.Record.Builder (build) as Record
import Data.Record.Builder (insert)
import Data.Record.Extra (eqRecord)
import Data.Symbol (SProxy(SProxy))
import Data.Traversable (for_, traverse, traverse_)
import Data.Tuple (Tuple(Tuple))
import Data.Variant (Variant, case_, inj)
import Data.Variant as V
import Genetics.Browser.Track.Backend (Peak, RenderedTrack, pixelSegments)
import Genetics.Browser.Track.Demo (Annotation, AnnotationField, SNP, annotationsForScale, demoBrowser, filterSig, getAnnotations, getGenes, getSNPs, showAnnotationField)
import Genetics.Browser.Track.UI.Canvas (BrowserCanvas, TrackPadding, _Dimensions, _Track, browserCanvas, browserOnClick, debugBrowserCanvas, dragScroll, renderBrowser, setBrowserCanvasSize, setElementStyle, wheelZoom)
import Genetics.Browser.Types (ChrId(ChrId), _NegLog10)
import Genetics.Browser.Types.Coordinates (CoordSys, CoordSysView(CoordSysView), _TotalSize, aroundPair, coordSys, normalizeView, pairSize, pairsOverlap, pixelsView, scaleViewBy, translateViewBy, viewScale, xPerPixel)
import Global.Unsafe (unsafeStringify)
import Graphics.Canvas as Canvas
import Graphics.Drawing (Drawing, Point)
import Math as Math
import Partial.Unsafe (unsafePartial)
import Simple.JSON (read)
import Unsafe.Coerce (unsafeCoerce)

foreign import windowInnerSize :: ∀ e. Eff e Canvas.Dimensions

-- | Set an event to fire on the given button id
foreign import buttonEvent :: ∀ e.
                              String
                           -> Eff e Unit
                           -> Eff e Unit

-- | Set callback to run after the window has resized (see UI.js for
-- | time waited), providing it with the new window size.
foreign import resizeEvent :: ∀ e.
                              ({ width  :: Number
                               , height :: Number } -> Eff e Unit)
                           -> Eff e Unit

-- | The value returned by the function that initializes the genome browser;
-- | provides an interface to the running browser instance.
type BrowserInterface e a
  = { getView          :: Aff e CoordSysView
    , getBrowserCanvas :: Aff e BrowserCanvas
    , lastOverlaps     :: Aff e (Number -> Point -> a)
    , queueCommand     :: Variant UICmdR -> Aff e Unit
    , queueUpdateView  :: UpdateView -> Eff e Unit
    }

-- | Creates the browser using the provided initial data, returning
-- | a BrowserInterface for reading state & sending commands to it
initializeBrowser :: CoordSys _ _
                  -> (BrowserCanvas -> CoordSysView
                      -> { tracks :: { snps :: RenderedTrack (SNP ())
                                     , annotations :: RenderedTrack (Annotation ()) }
                         , relativeUI :: Drawing
                         , fixedUI :: Drawing })
                  -> CoordSysView
                  -> BrowserCanvas
                  -> Aff _ (BrowserInterface _ ({snps :: Array (SNP ())}))
initializeBrowser cSys mainBrowser initView bc = do

  renderFiberVar <- AVar.makeEmptyVar

  viewVar <- AVar.makeVar initView
  bcVar <- AVar.makeVar bc

  uiCmdVar <- AVar.makeEmptyVar
  lastOverlapsVar <- AVar.makeEmptyVar

  drawCachedBrowser <- browserCache mainBrowser

  let getView = AVar.readVar viewVar
      getBrowserCanvas = AVar.readVar bcVar
      queueCommand = flip AVar.putVar uiCmdVar
      lastOverlaps = AVar.readVar lastOverlapsVar

    -- hardcoded timeout for now
  queueUpdateView <- uiViewUpdate cSys (wrap 100.0) {view: viewVar, uiCmd: uiCmdVar}

  let mainLoop = do
        uiCmd <- AVar.takeVar uiCmdVar
        case_ # V.on _render (\_ -> pure unit)
              # V.on _docResize (\ {width} -> do
                      canvas <- AVar.takeVar bcVar
                      let {height} = canvas ^. _Dimensions
                      canvas' <- liftEff $ setBrowserCanvasSize {width, height} canvas
                      AVar.putVar canvas' bcVar
                      queueCommand $ inj _render unit
                      mainLoop)
              $ uiCmd

        -- if there's a rendering fiber running, we kill it
        traverse_ (killFiber (error "Resetting renderer"))
          =<< AVar.tryTakeVar renderFiberVar

        csView <- getView
        canvas <- getBrowserCanvas
        toRender <- drawCachedBrowser canvas csView

        let currentScale = viewScale (canvas ^. _Track <<< _Dimensions) csView
            (Pair offset _) = pixelsView currentScale csView

                          -- offset the click by the current canvas translation
        _ <- AVar.tryTakeVar lastOverlapsVar
        AVar.putVar (\n r -> let r' = {x: r.x + offset, y: r.y }
                             in { snps: toRender.tracks.snps.overlaps n r' }) lastOverlapsVar

        -- fork a new renderFiber
        renderFiber <- forkAff
                       $ renderBrowser (wrap 2.0) canvas offset toRender
        AVar.putVar renderFiber renderFiberVar

        mainLoop

  _ <- forkAff mainLoop

  queueCommand $ inj _render unit

  pure { getView
       , getBrowserCanvas
       , lastOverlaps
       , queueCommand
       , queueUpdateView }

data UpdateView =
    ScrollView Number
  | ZoomView Number
  | ModView (Pair BigInt -> Pair BigInt)

instance showUpdateView :: Show UpdateView where
  show (ScrollView x) = "(Scroll by " <> show x <> ")"
  show (ZoomView s) = "(Zoom by " <> show s <> ")"
  show _ = "(ModView)"

    -- TODO idk if this instance makes sense??? whatevs
instance semigroupUpdateView :: Semigroup UpdateView where
  append (ScrollView x1) (ScrollView x2) = ScrollView (x1 + x2)
  append (ZoomView s1)   (ZoomView s2)   = ZoomView   (s1 * s2)
  append _ y = y

instance monoidUpdateView :: Monoid UpdateView where
  mempty = ModView id

updateViewFold :: UpdateView
               -> CoordSysView
               -> CoordSysView
updateViewFold uv iv = case uv of
  ZoomView   x -> iv `scaleViewBy`     x
  ScrollView x -> iv `translateViewBy` x
  ModView f    -> over CoordSysView f iv

queueCmd :: ∀ a. AVar a -> a -> Eff _ Unit
queueCmd av cmd = launchAff_ $ AVar.putVar cmd av

-- | Provided with the AVar containing the view state (which is updated atomically)
-- | and the UICmd AVar queue (which is only written to),
-- | return a function that queues view updates using the provided timeout.
uiViewUpdate :: ∀ r.
                CoordSys _ BigInt
             -> Milliseconds
             -> { view  :: AVar CoordSysView
                , uiCmd :: AVar (Variant UICmdR) | r }
             -> Aff _ (UpdateView -> Eff _ Unit)
uiViewUpdate cs timeout { view, uiCmd } = do
  viewCmd <- AVar.makeEmptyVar

  curCmdVar <- AVar.makeVar (ModView id)

  let normView = normalizeView cs (BigInt.fromInt 200000)
      loop' updater = do
        cmd <- AVar.takeVar viewCmd

        killFiber (error "Resetting view update") updater
        liftEff $ log $ "forking view update"

        curCmd <- AVar.takeVar curCmdVar
        let cmd' = curCmd <> cmd
        AVar.putVar cmd' curCmdVar

        updater' <- forkAff do
            delay timeout
            liftEff $ log "Running view update"

            vr <- AVar.takeVar view
            let vr' = normView $ updateViewFold cmd' vr
            AVar.putVar vr' view

            AVar.putVar (inj _render unit) uiCmd
            AVar.takeVar curCmdVar *> AVar.putVar mempty curCmdVar

        loop' updater'

  _ <- forkAff $ loop' (pure unit)

  pure $ queueCmd viewCmd

btnUI :: { scrollMod :: Number, zoomMod :: Number }
      -> (UpdateView -> Eff _ Unit)
      -> Eff _ Unit
btnUI mods cb = do
  buttonEvent "scrollLeft"  $ cb $ ScrollView     (-mods.scrollMod)
  buttonEvent "scrollRight" $ cb $ ScrollView       mods.scrollMod
  buttonEvent "zoomOut"     $ cb $ ZoomView $ 1.0 + mods.zoomMod
  buttonEvent "zoomIn"      $ cb $ ZoomView $ 1.0 - mods.zoomMod

type UICmdR = ( render :: Unit
              , docResize :: { width :: Number, height :: Number } )

_render = SProxy :: SProxy "render"
_docResize = SProxy :: SProxy "docResize"

debugView :: BrowserInterface _ _
          -> Eff _ { get :: _
                   , set :: _ }
debugView s = do
  let get name = launchAff_ do
         view <- s.getView
         liftEff do
           log $ "CoordSysView: " <> show (map BigInt.toString $ unwrap view)
           setWindow name $ (\(Pair l r) -> {l,r}) $ unwrap view

  let set lr = s.queueUpdateView
               $ ModView $ const (BigInt.fromNumber <$> Pair lr.l lr.r)

  pure {get, set}

browserCache
  :: ∀ a b.
     (BrowserCanvas
      -> CoordSysView
      -> { tracks :: { snps :: a
                     , annotations :: b }
         , relativeUI :: Drawing
         , fixedUI :: Drawing })
     -> Aff _ (BrowserCanvas
               -> CoordSysView
               -> Aff _ ({ tracks :: { snps :: a
                                     , annotations :: b }
                         , relativeUI :: Drawing
                         , fixedUI :: Drawing }))
browserCache f = do
  bcDim <- AVar.makeEmptyVar
  viewSize <- AVar.makeEmptyVar

  lastPartial <- AVar.makeEmptyVar
  lastFinal <- AVar.makeEmptyVar

  pure \bc csv -> do

    let newBCDim = bc ^. _Dimensions
    oldBCDim <- AVar.tryTakeVar bcDim
    AVar.putVar newBCDim bcDim

    oldPartial <- AVar.tryTakeVar lastPartial

    partial <- case oldBCDim, oldPartial of
      Just d, Just o -> do
        let changed = not $ d `eqRecord` newBCDim
        liftEff $ log $ "BrowserCanvas changed? " <> (show changed)
        -- remove the cached final output value on canvas resize,
        -- since we need to recalculate everything
        when changed (void $ AVar.tryTakeVar lastFinal)
        pure $ if d `eqRecord` newBCDim then o else f bc

      _, _ -> do
        liftEff $ log "Cache was empty! Recalculating all"
        _ <- AVar.tryTakeVar lastFinal
        pure $ f bc

    AVar.putVar partial lastPartial

    let newViewSize = pairSize $ unwrap csv
    oldViewSize <- AVar.tryTakeVar viewSize
    AVar.putVar newViewSize viewSize

    oldFinal <- AVar.tryTakeVar lastFinal

    let output = case oldViewSize, oldFinal of
          Just vs, Just o -> if newViewSize == vs then o else partial csv
          _, _ -> partial csv

    AVar.putVar output lastFinal
    pure output

printSNPInfo :: ∀ r. Array (SNP r) -> Eff _ Unit
printSNPInfo fs = do
  let n = length fs :: Int
      m = 5
  log $ "showing " <> show m <> " out of " <> show n <> " clicked glyphs"
  for_ (Array.take m fs) (log <<< unsafeCoerce)

wrapWith :: String -> String -> String
wrapWith tag x =
  "<"<> tag <>">" <> x <> "</"<> tag <>">"

snpHTML :: ∀ r.
           SNP r
        -> String
snpHTML {position, feature} = wrapWith "div" contents
  where contents = foldMap (wrapWith "p")
            [ "SNP: "    <> feature.name
            , "Chr: "    <> show feature.chrId
            , "Pos: "    <> show (Pair.fst position)
            , "Score: "  <> show feature.score
            , "-log10: " <> show (unwrap $ (feature.score ^. _NegLog10))
            ]

-- | Given a function to transform the data in the annotation's "rest" field
-- | to text (or Nothing if the field should not be displayed), produce a
-- | function that generates HTML from annotations
annotationHTML :: (AnnotationField -> Maybe String)
                -> Annotation () -> String
annotationHTML disp {feature} = wrapWith "div" contents
  where url = fromMaybe "No URL"
              $ map (\a -> "URL: <a target='_blank' href='"
                           <> a <> "'>" <> a <> "") feature.url

        name = fromMaybe ("Annotated SNP: " <> feature.name)
                         (("Gene: " <> _) <$> feature.gene)

        showOther fv = fv.field <> ": " <> (unsafeCoerce fv.value)

        contents = foldMap (wrapWith "p")
          $ [ name
            , "Pos: " <> show (feature.pos)
            , url
            , "Other data: "
            ] <> (filterMap disp
                  $ Array.fromFoldable feature.rest)

-- | Shows all data in "rest" using the default showAnnotationField (which uses unsafeCoerce)
annotationHTMLAll :: Annotation () -> String
annotationHTMLAll =
  annotationHTML (pure <<< showAnnotationField)

-- | Example HTML generator that only shows the "anno" field
annotationHTMLAnnoOnly :: Annotation () -> String
annotationHTMLAnnoOnly = annotationHTML disp
  where disp fv
          | fv.field == "anno" = pure $ showAnnotationField fv
          | otherwise          = Nothing

foreign import initDebugDiv :: ∀ e. Number -> Eff e Unit
foreign import setDebugDivVisibility :: ∀ e. String -> Eff e Unit
foreign import setDebugDivPoint :: ∀ e. Point -> Eff e Unit

foreign import setElementContents :: ∀ e. Element -> String -> Eff e Unit

data InfoBoxF
  = IBoxShow
  | IBoxHide
  | IBoxSetY Int
  | IBoxSetX Int
  | IBoxSetContents String

derive instance genericInfoBoxF :: Generic InfoBoxF _

instance showInfoBoxF :: Show InfoBoxF where
  show = genericShow

updateInfoBox :: Element -> InfoBoxF -> Eff _ Unit
updateInfoBox el cmd =
  case cmd of
    IBoxShow ->
      setElementStyle el "visibility" "visible"
    IBoxHide ->
      setElementStyle el "visibility" "hidden"
    (IBoxSetX x)    ->
      setElementStyle el "left" $ show x <> "px"
    (IBoxSetY y)    ->
      setElementStyle el "top"  $ show y <> "px"
    (IBoxSetContents html) ->
      setElementContents el html

infoBoxId :: ElementId
infoBoxId = wrap "infoBox"

initInfoBox :: Eff _ (InfoBoxF -> Eff _ Unit)
initInfoBox = do
  doc <- map DOM.htmlDocumentToDocument
           $ DOM.document =<< DOM.window
  el <- DOM.createElement "div" doc

  DOM.setId infoBoxId el

  DOM.documentElement doc >>= case _ of
    Nothing -> throw "Couldn't find document body!"
    Just docBody -> void $ DOM.appendChild el docBody

  pure $ updateInfoBox el

runBrowser :: Conf -> BrowserCanvas -> Eff _ _
runBrowser config bc = launchAff $ do

  let clickRadius = 1.0

  liftEff $ initDebugDiv clickRadius

  cmdInfoBox <- liftEff $ initInfoBox

  let cSys :: CoordSys ChrId BigInt
      cSys = coordSys mouseChrSizes

  {-
  let slots = uiSlots bc

      legend = { width: slots.right.size.width
               , entries: foldMap annotationLegendTest trackData.annotations }

      vscale =
        Record.build (insert (SProxy :: SProxy "width") slots.left.size.width
                  >>> insert (SProxy :: SProxy "color") black)
                      $ config.score

      sigSnps = foldMap (filterSig config.score) trackData.snps

      mainBrowser :: _
      mainBrowser = drawBrowser cSys {legend, vscale}
                      { snps:        renderSNPs vscale
                      , annotations: renderAnnotation cSys sigSnps vscale }
                      { snps:        fromMaybe mempty trackData.snps
                      , annotations: fromMaybe mempty trackData.annotations }
-}
  trackData <-
    {genes:_, snps:_, annotations:_}
    <$> foldMap (getGenes       cSys) config.urls.genes
    <*> foldMap (getSNPs        cSys) config.urls.snps
    <*> foldMap (getAnnotations cSys) config.urls.annotations

  let mainBrowser = demoBrowser cSys config trackData

  browser <-
    initializeBrowser cSys mainBrowser (wrap $ Pair zero (cSys^._TotalSize)) bc

  liftEff do
    resizeEvent \d ->
      launchAff_ $ browser.queueCommand $ inj _docResize d

    btnUI { scrollMod: 0.05, zoomMod: 0.1 } browser.queueUpdateView

    dragScroll bc \ {x,y} ->
       when (Math.abs x >= one) $ launchAff_ do
         trackDims <- Lens.view (_Track <<< _Dimensions)
                      <$> browser.getBrowserCanvas

         liftEff $ browser.queueUpdateView
                 $ ScrollView $ (-x) / trackDims.width

    let scrollZoomScale = 0.06
    wheelZoom bc \dY ->
       browser.queueUpdateView
         $ ZoomView $ 1.0 + scrollZoomScale * dY

  liftEff do
    let overlayDebug :: _
        overlayDebug p = do
          setDebugDivVisibility "visible"
          setDebugDivPoint p

    let annotAround r snp =
          Array.find (\a -> ((r `aroundPair` a.position)
                             `pairsOverlap` snp.position))
            =<< Map.lookup snp.feature.chrId
            =<< trackData.annotations

        glyphClick :: _
        glyphClick p = launchAff_ do

          v <- browser.getView
          let vs = viewScale (bc ^. _Track <<< _Dimensions) v
              radius = wrap $ (xPerPixel vs) * 3.75

          lastOverlaps' <- browser.lastOverlaps
          liftEff do
            let clicked = (lastOverlaps' clickRadius p).snps
            printSNPInfo clicked
            case Array.head clicked of
              Nothing -> cmdInfoBox IBoxHide
              Just g  -> do
                cmdInfoBox IBoxShow
                cmdInfoBox $ IBoxSetX $ Int.round p.x
                cmdInfoBox $ IBoxSetContents
                  $ fromMaybe (snpHTML g)
                    (annotationHTMLAll <$> annotAround radius g)

    browserOnClick bc
      { overlay: \_ -> pure unit
      -- { overlay: overlayDebug
      , track:   glyphClick }

  pure unit

type DataURLs = { snps        :: Maybe String
                , annotations :: Maybe String
                , genes       :: Maybe String
                }

type Conf = { browserHeight :: Number
            , trackPadding :: TrackPadding
            , score :: { min :: Number, max :: Number, sig :: Number }
            , urls :: DataURLs
            }

foreign import setWindow :: ∀ e a. String -> a -> Eff e Unit

main :: Foreign -> Eff _ _
main rawConfig = do

  el' <- do
    doc <- DOM.htmlDocumentToDocument
           <$> (DOM.document =<< DOM.window)
    DOM.querySelector (wrap "#browser") (toParentNode doc)

  case read rawConfig :: Either MultipleErrors Conf of
    Left errs -> traverse_ (log <<< renderForeignError) errs
    Right c   -> do
      case el' of
        Nothing -> log "Could not find element '#browser'"
        Just el -> do

          {width} <- windowInnerSize
          let dimensions = { width, height: c.browserHeight }
          bc <- browserCanvas dimensions c.trackPadding el

          debugBrowserCanvas "debugBC" bc

          log $ unsafeStringify c
          void $ runBrowser c bc

-- TODO this could almost certainly be done better
chunkConsumer :: ∀ m a.
                 Foldable m
              => Monoid (m a)
              => AVar (m a)
              -> Consumer (m a) (Aff _) Unit
chunkConsumer av = Co.consumer \m ->
  if null m
    then pure $ Just unit
    else do
      sofar <- AVar.takeVar av
      AVar.putVar (sofar <> m) av
      delay (wrap 20.0)
      pure Nothing

type TrackVar a = AVar (Map ChrId (Array a))
type TrackProducer eff a = Producer (Map ChrId (Array a)) (Aff eff) Unit

-- Feels like this one takes care of a bit too much...
fetchLoop1 :: ∀ a.
              (Maybe (Aff _ (TrackProducer _ a)))
           -> Aff _ (TrackVar a)
fetchLoop1 Nothing = AVar.makeEmptyVar
fetchLoop1 (Just startProd) = do
  prod <- startProd
  avar <- AVar.makeVar mempty
  _ <- forkAff $ runProcess $ prod `connect` chunkConsumer avar
  pure avar

-- | Starts threads that fetch & parse each of the provided tracks,
-- | filling an AVar over time per track, which can be used by other parts of the application
-- | (read only, should be a newtype)
{-
fetchLoop :: CoordSys ChrId BigInt
          -> DataURLs
          -> Aff _
               { gwas :: TrackVar (GWASFeature ())
               , genes :: TrackVar BedFeature
               }
fetchLoop cs urls = do
  gwas <-        fetchLoop1 $ produceGWAS   cs <$> urls.gwas
  genes <-       fetchLoop1 $ produceGenes  cs <$> urls.genes
  pure { gwas, genes }
-}

mouseChrSizes :: Array (Tuple ChrId BigInt)
mouseChrSizes =
  unsafePartial
  $ map (bimap ChrId (fromJust <<< BigInt.fromString))
      [ Tuple "1"   "195471971"
      , Tuple "2"   "182113224"
      , Tuple "3"   "160039680"
      , Tuple "4"   "156508116"
      , Tuple "5"   "151834684"
      , Tuple "6"   "149736546"
      , Tuple "7"   "145441459"
      , Tuple "8"   "129401213"
      , Tuple "9"   "124595110"
      , Tuple "10"  "130694993"
      , Tuple "11"  "122082543"
      , Tuple "12"  "120129022"
      , Tuple "13"  "120421639"
      , Tuple "14"  "124902244"
      , Tuple "15"  "104043685"
      , Tuple "16"  "98207768"
      , Tuple "17"  "94987271"
      , Tuple "18"  "90702639"
      , Tuple "19"  "61431566"
      , Tuple "X"   "171031299"
      , Tuple "Y"   "91744698"
      ]


Update for changes in strings

Update for changes in strings

module Test.QuickCheck.Arbitrary
  ( class Arbitrary
  , arbitrary
  , class Coarbitrary
  , coarbitrary
  , genericArbitrary
  , genericCoarbitrary
  , class ArbitraryGenericSum
  , arbitraryGenericSum
  , class ArbitraryRowList
  , arbitraryRecord
  ) where

import Prelude

import Control.Monad.Gen.Class (chooseBool)
import Control.Monad.Gen.Common as MGC
import Control.Monad.ST as ST
import Data.Array.NonEmpty (NonEmptyArray)
import Data.Array.NonEmpty as NEA
import Data.Array.ST as STA
import Data.Either (Either(..))
import Data.Enum (fromEnum, toEnumWithDefaults)
import Data.Foldable (foldl)
import Data.Generic.Rep (class Generic, to, from, NoArguments(..), Sum(..), Product(..), Constructor(..), Argument(..))
import Data.Identity (Identity(..))
import Data.Int (toNumber)
import Data.Lazy (Lazy, defer, force)
import Data.List (List)
import Data.List.NonEmpty (NonEmptyList(..))
import Data.Maybe (Maybe(..), fromJust)
import Data.Newtype (wrap)
import Data.NonEmpty (NonEmpty(..), (:|))
import Data.String.CodeUnits (charAt, fromCharArray, split)
import Data.String.NonEmpty (NonEmptyString)
import Data.String.NonEmpty.CodeUnits as NES
import Data.Symbol (class IsSymbol, SProxy(..))
import Data.Tuple (Tuple(..))
import Partial.Unsafe (unsafePartial)
import Prim.Row as Row
import Prim.RowList as RL
import Record as Record
import Test.QuickCheck.Gen (Gen, arrayOf, chooseInt, elements, listOf, oneOf, perturbGen, repeatable, sized, uniform)
import Type.Data.RowList (RLProxy(..))

-- | The `Arbitrary` class represents those types whose values can be
-- | _randomly-generated_.
-- |
-- | `arbitrary` uses the `Gen` monad to express a random generator for
-- | the type `t`. Combinators in the `Test.QuickCheck.Gen`
-- | module can be used to construct random generators.
class Arbitrary t where
  arbitrary :: Gen t

-- | The `Coarbitrary` class represents types which appear on the left of
-- | an `Arbitrary` function arrow.
-- |
-- | To construct an `Arbitrary` instance for the type `a -> b`, we need to
-- | use the input of type `a` to _perturb_ a random generator for `b`. This
-- | is the role of the `coarbitrary` function.
-- |
-- | `Coarbitrary` instances can be written using the `perturbGen` function.
class Coarbitrary t where
  coarbitrary :: forall r. t -> Gen r -> Gen r

instance arbBoolean :: Arbitrary Boolean where
  arbitrary = chooseBool

instance coarbBoolean :: Coarbitrary Boolean where
  coarbitrary true = perturbGen 1.0
  coarbitrary false = perturbGen 2.0

instance arbNumber :: Arbitrary Number where
  arbitrary = uniform

instance coarbNumber :: Coarbitrary Number where
  coarbitrary = perturbGen

instance arbInt :: Arbitrary Int where
  arbitrary = chooseInt (-1000000) 1000000

instance coarbInt :: Coarbitrary Int where
  coarbitrary = perturbGen <<< toNumber

instance arbString :: Arbitrary String where
  arbitrary = fromCharArray <$> arbitrary

instance coarbString :: Coarbitrary String where
  coarbitrary s = coarbitrary $ (charAt zero <$> split (wrap "") s)

instance arbNonEmptyString :: Arbitrary NonEmptyString where
  arbitrary = NES.cons <$> arbitrary <*> arbitrary

instance coarbNonEmptyString :: Coarbitrary NonEmptyString where
  coarbitrary = coarbitrary <<< NES.toString

instance arbChar :: Arbitrary Char where
  arbitrary = toEnumWithDefaults bottom top <$> chooseInt 0 65536

instance coarbChar :: Coarbitrary Char where
  coarbitrary c = coarbitrary $ fromEnum c

instance arbUnit :: Arbitrary Unit where
  arbitrary = pure unit

instance coarbUnit :: Coarbitrary Unit where
  coarbitrary _ = perturbGen 1.0

instance arbOrdering :: Arbitrary Ordering where
  arbitrary = elements $ LT :| [EQ, GT]

instance coarbOrdering :: Coarbitrary Ordering where
  coarbitrary LT = perturbGen 1.0
  coarbitrary EQ = perturbGen 2.0
  coarbitrary GT = perturbGen 3.0

instance arbArray :: Arbitrary a => Arbitrary (Array a) where
  arbitrary = arrayOf arbitrary

instance coarbArray :: Coarbitrary a => Coarbitrary (Array a) where
  coarbitrary = foldl (\f x -> f <<< coarbitrary x) identity

instance arbNonEmptyArray :: Arbitrary a => Arbitrary (NonEmptyArray a) where
  arbitrary = do
     x <- arbitrary
     xs <- arbitrary
     pure $ unsafePartial fromJust $ NEA.fromArray $ ST.run do
        mxs <- STA.unsafeThaw xs
        _ <- STA.push x mxs
        STA.unsafeFreeze mxs

instance coarbNonEmptyArray :: Coarbitrary a => Coarbitrary (NonEmptyArray a) where
  coarbitrary = coarbitrary <<< NEA.toArray

instance arbFunction :: (Coarbitrary a, Arbitrary b) => Arbitrary (a -> b) where
  arbitrary = repeatable (\a -> coarbitrary a arbitrary)

instance coarbFunction :: (Arbitrary a, Coarbitrary b) => Coarbitrary (a -> b) where
  coarbitrary f gen = do
    xs <- arbitrary
    coarbitrary (map f (xs :: Array a)) gen

instance arbTuple :: (Arbitrary a, Arbitrary b) => Arbitrary (Tuple a b) where
  arbitrary = Tuple <$> arbitrary <*> arbitrary

instance coarbTuple :: (Coarbitrary a, Coarbitrary b) => Coarbitrary (Tuple a b) where
  coarbitrary (Tuple a b) = coarbitrary a >>> coarbitrary b

instance arbMaybe :: Arbitrary a => Arbitrary (Maybe a) where
  arbitrary = MGC.genMaybe arbitrary

instance coarbMaybe :: Coarbitrary a => Coarbitrary (Maybe a) where
  coarbitrary Nothing = perturbGen 1.0
  coarbitrary (Just a) = coarbitrary a

instance arbEither :: (Arbitrary a, Arbitrary b) => Arbitrary (Either a b) where
  arbitrary = MGC.genEither arbitrary arbitrary

instance coarbEither :: (Coarbitrary a, Coarbitrary b) => Coarbitrary (Either a b) where
  coarbitrary (Left a)  = coarbitrary a
  coarbitrary (Right b) = coarbitrary b

instance arbitraryList :: Arbitrary a => Arbitrary (List a) where
  arbitrary = sized \n -> chooseInt zero n >>= flip listOf arbitrary

instance coarbList :: Coarbitrary a => Coarbitrary (List a) where
  coarbitrary = foldl (\f x -> f <<< coarbitrary x) identity

instance arbitraryIdentity :: Arbitrary a => Arbitrary (Identity a) where
  arbitrary = Identity <$> arbitrary

instance coarbIdentity :: Coarbitrary a => Coarbitrary (Identity a) where
  coarbitrary (Identity a) = coarbitrary a

instance arbitraryLazy :: Arbitrary a => Arbitrary (Lazy a) where
  arbitrary = arbitrary >>= pure <<< defer <<< const

instance coarbLazy :: Coarbitrary a => Coarbitrary (Lazy a) where
  coarbitrary a = coarbitrary (force a)

instance arbNonEmpty :: (Arbitrary (f a), Arbitrary a) => Arbitrary (NonEmpty f a) where
  arbitrary = NonEmpty <$> arbitrary <*> arbitrary

instance coarbNonEmpty :: (Coarbitrary (f a), Coarbitrary a) => Coarbitrary (NonEmpty f a) where
  coarbitrary (NonEmpty head tail) = coarbitrary head >>> coarbitrary tail

instance arbNonEmptyList :: Arbitrary a => Arbitrary (NonEmptyList a) where
  arbitrary = NonEmptyList <$> arbitrary

instance coarbNonEmptyList :: Coarbitrary a => Coarbitrary (NonEmptyList a) where
  coarbitrary (NonEmptyList nel) = coarbitrary nel

instance arbitraryNoArguments :: Arbitrary NoArguments where
  arbitrary = pure NoArguments

instance coarbitraryNoArguments :: Coarbitrary NoArguments where
  coarbitrary NoArguments = identity

-- | To be able to evenly distribute over chains of Sum types we build up
-- | a collection of generators and choose between.  Each right component
-- | of a Sum is either a Constructor or another Sum.
class ArbitraryGenericSum t where
  arbitraryGenericSum :: Array (Gen t)

instance arbGenSumSum :: (Arbitrary l, ArbitraryGenericSum r) => ArbitraryGenericSum (Sum l r) where
  arbitraryGenericSum = [Inl <$> arbitrary] <> (map Inr <$> arbitraryGenericSum)

instance arbGenSumConstructor :: Arbitrary a => ArbitraryGenericSum (Constructor s a) where
  arbitraryGenericSum = [arbitrary]

instance arbitrarySum :: (Arbitrary l, ArbitraryGenericSum r) => Arbitrary (Sum l r) where
  arbitrary = oneOf $ (Inl <$> arbitrary) :| (map Inr <$> arbitraryGenericSum)

instance coarbitrarySum :: (Coarbitrary l, Coarbitrary r) => Coarbitrary (Sum l r) where
  coarbitrary (Inl l) = coarbitrary l
  coarbitrary (Inr r) = coarbitrary r

instance arbitraryProduct :: (Arbitrary l, Arbitrary r) => Arbitrary (Product l r) where
  arbitrary = Product <$> arbitrary <*> arbitrary

instance coarbitraryProduct :: (Coarbitrary l, Coarbitrary r) => Coarbitrary (Product l r) where
  coarbitrary (Product a b) = coarbitrary a >>> coarbitrary b

instance arbitraryConstructor :: Arbitrary a => Arbitrary (Constructor s a) where
  arbitrary = Constructor <$> arbitrary

instance coarbitraryConstructor :: Coarbitrary a => Coarbitrary (Constructor s a) where
  coarbitrary (Constructor a) = coarbitrary a

instance arbitraryArgument :: Arbitrary a => Arbitrary (Argument a) where
  arbitrary = Argument <$> arbitrary

instance coarbitraryArgument :: Coarbitrary a => Coarbitrary (Argument a) where
  coarbitrary (Argument a) = coarbitrary a

-- | A `Generic` implementation of the `arbitrary` member from the `Arbitrary` type class.
genericArbitrary :: forall a rep. Generic a rep => Arbitrary rep => Gen a
genericArbitrary = to <$> (arbitrary :: Gen rep)

-- | A `Generic` implementation of the `coarbitrary` member from the `Coarbitrary` type class.
genericCoarbitrary :: forall a rep. Generic a rep => Coarbitrary rep => a -> Gen a -> Gen a
genericCoarbitrary x g = to <$> coarbitrary (from x) (from <$> g)

-- | A helper typeclass to implement `Arbitrary` for records.
class ArbitraryRowList list row | list -> row where
  arbitraryRecord :: RLProxy list -> Gen (Record row)

instance arbitraryRowListNil :: ArbitraryRowList RL.Nil () where
  arbitraryRecord _ = pure {}

instance arbitraryRowListCons ::
  ( Arbitrary a
  , ArbitraryRowList listRest rowRest
  , Row.Lacks key rowRest
  , Row.Cons key a rowRest rowFull
  , RL.RowToList rowFull (RL.Cons key a listRest)
  , IsSymbol key
  ) => ArbitraryRowList (RL.Cons key a listRest) rowFull where
  arbitraryRecord _ = do
     value <- arbitrary
     previous <- arbitraryRecord (RLProxy :: RLProxy listRest)
     pure $ Record.insert (SProxy :: SProxy key) value previous

instance arbitraryRecordInstance ::
  ( RL.RowToList row list
  , ArbitraryRowList list row
  ) => Arbitrary (Record row) where
  arbitrary = arbitraryRecord (RLProxy :: RLProxy list)

module Test.QuickCheck.Arbitrary
  ( class Arbitrary
  , arbitrary
  , class Coarbitrary
  , coarbitrary
  , genericArbitrary
  , genericCoarbitrary
  , class ArbitraryGenericSum
  , arbitraryGenericSum
  , class ArbitraryRowList
  , arbitraryRecord
  ) where

import Prelude

import Control.Monad.Gen.Class (chooseBool)
import Control.Monad.Gen.Common as MGC
import Control.Monad.ST as ST
import Data.Array.NonEmpty (NonEmptyArray)
import Data.Array.NonEmpty as NEA
import Data.Array.ST as STA
import Data.Char (toCharCode, fromCharCode)
import Data.Either (Either(..))
import Data.Foldable (foldl)
import Data.Generic.Rep (class Generic, to, from, NoArguments(..), Sum(..), Product(..), Constructor(..), Argument(..))
import Data.Identity (Identity(..))
import Data.Int (toNumber)
import Data.Lazy (Lazy, defer, force)
import Data.List (List)
import Data.List.NonEmpty (NonEmptyList(..))
import Data.Maybe (Maybe(..), fromJust)
import Data.Newtype (wrap)
import Data.NonEmpty (NonEmpty(..), (:|))
import Record as Record
import Data.String (charCodeAt, fromCharArray, split)
import Data.String.NonEmpty (NonEmptyString)
import Data.String.NonEmpty as NES
import Data.Symbol (class IsSymbol, SProxy(..))
import Data.Tuple (Tuple(..))
import Partial.Unsafe (unsafePartial)
import Prim.Row as Row
import Prim.RowList as RL
import Test.QuickCheck.Gen (Gen, arrayOf, chooseInt, elements, listOf, oneOf, perturbGen, repeatable, sized, uniform)
import Type.Data.RowList (RLProxy(..))

-- | The `Arbitrary` class represents those types whose values can be
-- | _randomly-generated_.
-- |
-- | `arbitrary` uses the `Gen` monad to express a random generator for
-- | the type `t`. Combinators in the `Test.QuickCheck.Gen`
-- | module can be used to construct random generators.
class Arbitrary t where
  arbitrary :: Gen t

-- | The `Coarbitrary` class represents types which appear on the left of
-- | an `Arbitrary` function arrow.
-- |
-- | To construct an `Arbitrary` instance for the type `a -> b`, we need to
-- | use the input of type `a` to _perturb_ a random generator for `b`. This
-- | is the role of the `coarbitrary` function.
-- |
-- | `Coarbitrary` instances can be written using the `perturbGen` function.
class Coarbitrary t where
  coarbitrary :: forall r. t -> Gen r -> Gen r

instance arbBoolean :: Arbitrary Boolean where
  arbitrary = chooseBool

instance coarbBoolean :: Coarbitrary Boolean where
  coarbitrary true = perturbGen 1.0
  coarbitrary false = perturbGen 2.0

instance arbNumber :: Arbitrary Number where
  arbitrary = uniform

instance coarbNumber :: Coarbitrary Number where
  coarbitrary = perturbGen

instance arbInt :: Arbitrary Int where
  arbitrary = chooseInt (-1000000) 1000000

instance coarbInt :: Coarbitrary Int where
  coarbitrary = perturbGen <<< toNumber

instance arbString :: Arbitrary String where
  arbitrary = fromCharArray <$> arbitrary

instance coarbString :: Coarbitrary String where
  coarbitrary s = coarbitrary $ (charCodeAt zero <$> split (wrap "") s)

instance arbNonEmptyString :: Arbitrary NonEmptyString where
  arbitrary = NES.cons <$> arbitrary <*> arbitrary

instance coarbNonEmptyString :: Coarbitrary NonEmptyString where
  coarbitrary = coarbitrary <<< NES.toString

instance arbChar :: Arbitrary Char where
  arbitrary = fromCharCode <$> chooseInt 0 65536

instance coarbChar :: Coarbitrary Char where
  coarbitrary c = coarbitrary $ toCharCode c

instance arbUnit :: Arbitrary Unit where
  arbitrary = pure unit

instance coarbUnit :: Coarbitrary Unit where
  coarbitrary _ = perturbGen 1.0

instance arbOrdering :: Arbitrary Ordering where
  arbitrary = elements $ LT :| [EQ, GT]

instance coarbOrdering :: Coarbitrary Ordering where
  coarbitrary LT = perturbGen 1.0
  coarbitrary EQ = perturbGen 2.0
  coarbitrary GT = perturbGen 3.0

instance arbArray :: Arbitrary a => Arbitrary (Array a) where
  arbitrary = arrayOf arbitrary

instance coarbArray :: Coarbitrary a => Coarbitrary (Array a) where
  coarbitrary = foldl (\f x -> f <<< coarbitrary x) identity

instance arbNonEmptyArray :: Arbitrary a => Arbitrary (NonEmptyArray a) where
  arbitrary = do
     x <- arbitrary
     xs <- arbitrary
     pure $ unsafePartial fromJust $ NEA.fromArray $ ST.run do
        mxs <- STA.unsafeThaw xs
        _ <- STA.push x mxs
        STA.unsafeFreeze mxs

instance coarbNonEmptyArray :: Coarbitrary a => Coarbitrary (NonEmptyArray a) where
  coarbitrary = coarbitrary <<< NEA.toArray

instance arbFunction :: (Coarbitrary a, Arbitrary b) => Arbitrary (a -> b) where
  arbitrary = repeatable (\a -> coarbitrary a arbitrary)

instance coarbFunction :: (Arbitrary a, Coarbitrary b) => Coarbitrary (a -> b) where
  coarbitrary f gen = do
    xs <- arbitrary
    coarbitrary (map f (xs :: Array a)) gen

instance arbTuple :: (Arbitrary a, Arbitrary b) => Arbitrary (Tuple a b) where
  arbitrary = Tuple <$> arbitrary <*> arbitrary

instance coarbTuple :: (Coarbitrary a, Coarbitrary b) => Coarbitrary (Tuple a b) where
  coarbitrary (Tuple a b) = coarbitrary a >>> coarbitrary b

instance arbMaybe :: Arbitrary a => Arbitrary (Maybe a) where
  arbitrary = MGC.genMaybe arbitrary

instance coarbMaybe :: Coarbitrary a => Coarbitrary (Maybe a) where
  coarbitrary Nothing = perturbGen 1.0
  coarbitrary (Just a) = coarbitrary a

instance arbEither :: (Arbitrary a, Arbitrary b) => Arbitrary (Either a b) where
  arbitrary = MGC.genEither arbitrary arbitrary

instance coarbEither :: (Coarbitrary a, Coarbitrary b) => Coarbitrary (Either a b) where
  coarbitrary (Left a)  = coarbitrary a
  coarbitrary (Right b) = coarbitrary b

instance arbitraryList :: Arbitrary a => Arbitrary (List a) where
  arbitrary = sized \n -> chooseInt zero n >>= flip listOf arbitrary

instance coarbList :: Coarbitrary a => Coarbitrary (List a) where
  coarbitrary = foldl (\f x -> f <<< coarbitrary x) identity

instance arbitraryIdentity :: Arbitrary a => Arbitrary (Identity a) where
  arbitrary = Identity <$> arbitrary

instance coarbIdentity :: Coarbitrary a => Coarbitrary (Identity a) where
  coarbitrary (Identity a) = coarbitrary a

instance arbitraryLazy :: Arbitrary a => Arbitrary (Lazy a) where
  arbitrary = arbitrary >>= pure <<< defer <<< const

instance coarbLazy :: Coarbitrary a => Coarbitrary (Lazy a) where
  coarbitrary a = coarbitrary (force a)

instance arbNonEmpty :: (Arbitrary (f a), Arbitrary a) => Arbitrary (NonEmpty f a) where
  arbitrary = NonEmpty <$> arbitrary <*> arbitrary

instance coarbNonEmpty :: (Coarbitrary (f a), Coarbitrary a) => Coarbitrary (NonEmpty f a) where
  coarbitrary (NonEmpty head tail) = coarbitrary head >>> coarbitrary tail

instance arbNonEmptyList :: Arbitrary a => Arbitrary (NonEmptyList a) where
  arbitrary = NonEmptyList <$> arbitrary

instance coarbNonEmptyList :: Coarbitrary a => Coarbitrary (NonEmptyList a) where
  coarbitrary (NonEmptyList nel) = coarbitrary nel

instance arbitraryNoArguments :: Arbitrary NoArguments where
  arbitrary = pure NoArguments

instance coarbitraryNoArguments :: Coarbitrary NoArguments where
  coarbitrary NoArguments = identity

-- | To be able to evenly distribute over chains of Sum types we build up
-- | a collection of generators and choose between.  Each right component
-- | of a Sum is either a Constructor or another Sum.
class ArbitraryGenericSum t where
  arbitraryGenericSum :: Array (Gen t)

instance arbGenSumSum :: (Arbitrary l, ArbitraryGenericSum r) => ArbitraryGenericSum (Sum l r) where
  arbitraryGenericSum = [Inl <$> arbitrary] <> (map Inr <$> arbitraryGenericSum)

instance arbGenSumConstructor :: Arbitrary a => ArbitraryGenericSum (Constructor s a) where
  arbitraryGenericSum = [arbitrary]

instance arbitrarySum :: (Arbitrary l, ArbitraryGenericSum r) => Arbitrary (Sum l r) where
  arbitrary = oneOf $ (Inl <$> arbitrary) :| (map Inr <$> arbitraryGenericSum)

instance coarbitrarySum :: (Coarbitrary l, Coarbitrary r) => Coarbitrary (Sum l r) where
  coarbitrary (Inl l) = coarbitrary l
  coarbitrary (Inr r) = coarbitrary r

instance arbitraryProduct :: (Arbitrary l, Arbitrary r) => Arbitrary (Product l r) where
  arbitrary = Product <$> arbitrary <*> arbitrary

instance coarbitraryProduct :: (Coarbitrary l, Coarbitrary r) => Coarbitrary (Product l r) where
  coarbitrary (Product a b) = coarbitrary a >>> coarbitrary b

instance arbitraryConstructor :: Arbitrary a => Arbitrary (Constructor s a) where
  arbitrary = Constructor <$> arbitrary

instance coarbitraryConstructor :: Coarbitrary a => Coarbitrary (Constructor s a) where
  coarbitrary (Constructor a) = coarbitrary a

instance arbitraryArgument :: Arbitrary a => Arbitrary (Argument a) where
  arbitrary = Argument <$> arbitrary

instance coarbitraryArgument :: Coarbitrary a => Coarbitrary (Argument a) where
  coarbitrary (Argument a) = coarbitrary a

-- | A `Generic` implementation of the `arbitrary` member from the `Arbitrary` type class.
genericArbitrary :: forall a rep. Generic a rep => Arbitrary rep => Gen a
genericArbitrary = to <$> (arbitrary :: Gen rep)

-- | A `Generic` implementation of the `coarbitrary` member from the `Coarbitrary` type class.
genericCoarbitrary :: forall a rep. Generic a rep => Coarbitrary rep => a -> Gen a -> Gen a
genericCoarbitrary x g = to <$> coarbitrary (from x) (from <$> g)

-- | A helper typeclass to implement `Arbitrary` for records.
class ArbitraryRowList list row | list -> row where
  arbitraryRecord :: RLProxy list -> Gen (Record row)

instance arbitraryRowListNil :: ArbitraryRowList RL.Nil () where
  arbitraryRecord _ = pure {}

instance arbitraryRowListCons ::
  ( Arbitrary a
  , ArbitraryRowList listRest rowRest
  , Row.Lacks key rowRest
  , Row.Cons key a rowRest rowFull
  , RL.RowToList rowFull (RL.Cons key a listRest)
  , IsSymbol key
  ) => ArbitraryRowList (RL.Cons key a listRest) rowFull where
  arbitraryRecord _ = do
     value <- arbitrary
     previous <- arbitraryRecord (RLProxy :: RLProxy listRest)
     pure $ Record.insert (SProxy :: SProxy key) value previous

instance arbitraryRecordInstance ::
  ( RL.RowToList row list
  , ArbitraryRowList list row
  ) => Arbitrary (Record row) where
  arbitrary = arbitraryRecord (RLProxy :: RLProxy list)


formatting

formatting

module Server where

import Debug.Trace
import Data.Tuple
import Data.JSON
import Data.Function
import Data.Maybe
import Data.Array
import qualified Data.String as S
import qualified Data.Either as E
import qualified Data.Map as M
import Data.Foldable (for_, all, find)
import Control.Monad
import Control.Monad.Eff
import Control.Monad.Eff.Ref
import Control.Reactive.Timer
import Control.Lens (lens, (^.), (%~), (.~), at, LensP())

import qualified NodeWebSocket as WS
import qualified NodeHttp as Http
import Types
import Game
import Utils

initialState :: ServerState
initialState =
  { gameState: WaitingForPlayers M.empty
  , connections: []
  }

stepsPerSecond = 30

main = do
  refState <- newRef initialState
  port <- portOrDefault 8080
  httpServer <- createHttpServer
  wsServer <- createWebSocketServer refState

  WS.mount wsServer httpServer
  Http.listen httpServer port
  trace $ "listening on " <> show port <> "..."

  void $ interval (1000 / stepsPerSecond) $
    runCallback refState step {}

createHttpServer =
  Http.createServer $ \req res -> do
    let path = (Http.getUrl req).pathname
    let reply =
      case path of
          "/"           -> Http.sendFile "static/index.html"
          "/js/game.js" -> Http.sendFile "static/js/game.js"
          _             -> Http.send404
    reply res

createWebSocketServer refState = do
  server <- WS.mkServer
  WS.onRequest server $ \req -> do
    trace "got a request"
    let playerName = S.drop 1 (WS.resourceUrl req).search

    if S.null playerName
      then WS.reject req
      else do
        conn <- WS.accept req
        maybePId <- tryAddPlayer conn refState playerName

        case maybePId of
          Just pId -> do
            trace $ "opened connection for player " <>
                        show pId <> ": " <> playerName
            runCallback refState onNewPlayer {pId:pId}

            WS.onMessage conn $ \msg ->
              runCallback refState onMessage {msg:msg, pId:pId}
            WS.onClose   conn $ \close ->
              runCallback refState onClose {pId:pId}
          Nothing -> do
            trace "rejecting connection, no player ids available"
            WS.close conn

  return server

runCallback refState callback args = do
  state <- readRef refState
  state' <- callback args state
  writeRef refState state'

tryAddPlayer conn refState name = do
  state <- readRef refState
  case getNextPlayerId state of
    Just pId -> do
      let state' = state { connections =
                            state.connections <>
                              [{ pId: pId, wsConn: conn, name: name }] }
      writeRef refState state'
      return (Just pId)
    Nothing ->
      return Nothing

getNextPlayerId state =
  let playerIdsInUse = map (\c -> c.pId) state.connections
  in  head $ allPlayerIds \\ playerIdsInUse

gameState = lens (\s -> s.gameState) (\s x -> s { gameState = x })

step args state = do
  case state ^. gameState of
    InProgress g -> do
      let r = stepGame g.input g.game
      let game' = fst r
      let updates = snd r
      let newState =
            if isEnded game'
              then WaitingForPlayers M.empty
              else InProgress { game: game', input: M.empty }

      sendUpdates state updates
      return $ state { gameState = newState }

    WaitingForPlayers m -> do
      if readyToStart m
        then do
          trace "all players are ready; starting game"
          let game = makeGame (M.keys m)
          sendUpdates state $ GameStarting game
          let newState = InProgress { game: game, input: M.empty }
          return $ state { gameState = newState }
        else
          return state
      where
      minPlayers = 2
      readyToStart m =
        let ps = M.values m
        in length ps >= minPlayers && all id ps

onMessage args state = do
  case state ^. gameState of
    InProgress g ->
      case (eitherDecode args.msg :: E.Either String Direction) of
        E.Right newDir -> do
          let newInput = M.insert args.pId (Just newDir) g.input
          return $ state
              { gameState = InProgress { game: g.game, input: newInput }}
        E.Left err -> do
          trace $ "failed to parse message: " <> err
          return state

    WaitingForPlayers m -> do
      case (eitherDecode args.msg :: E.Either String Boolean) of
         E.Right isReady -> do
           trace $ "updated ready state for " <> show args.pId <>
                     ": " <> show isReady
           let m' = M.insert args.pId isReady m
           return $ state { gameState = WaitingForPlayers m' }
         E.Left err -> do
           trace $ "failed to parse message: " <> err
           return state

onNewPlayer args state = do
  sendUpdateTo state args.pId (YourPlayerIdIs args.pId)
  return state

onClose args state = do
  let conns = filter (\c -> (args.pId::PlayerId) /= c.pId) state.connections
  let s' = state { connections = conns }
  trace $ "closed connection for " <> show args.pId
  return s'

sendUpdates :: forall e a. (ToJSON a) =>
  ServerState -> a -> Eff (ws :: WS.WebSocket | e) Unit
sendUpdates state updates = do
  for_ state.connections $ \c ->
    WS.send c.wsConn $ encode updates

sendUpdateTo :: forall e a. (ToJSON a) =>
  ServerState -> PlayerId -> a -> Eff (ws :: WS.WebSocket | e) Unit
sendUpdateTo state pId update = do
  let mConn = find (\c -> c.pId == pId) state.connections
  whenJust mConn $ \conn ->
    WS.send conn.wsConn $ encode update

foreign import chdir
  """
  function chdir(path) {
    return function() {
      process.chdir(path)
    }
  }
  """ :: forall e. String -> Eff (process :: Process | e) Unit

foreign import getEnvImpl
  """
  function getEnvImpl(just, nothing, key) {
    return function() {
      var v = process.env[key]
      return v ? just(v) : nothing
    }
  }
  """ :: forall e a.
  Fn3
    (a -> Maybe a)
    (Maybe a)
    String
    (Eff (process :: Process | e) (Maybe String))

getEnv :: forall e.
  String -> Eff (process :: Process | e) (Maybe String)
getEnv key =
  runFn3 getEnvImpl Just Nothing key

parseNumber :: String -> Maybe Number
parseNumber = decode

portOrDefault :: forall e.
  Number -> Eff (process :: Process | e) Number
portOrDefault default = do
  port <- getEnv "PORT"
  return $ fromMaybe default (port >>= parseNumber)