-- Skrepka Relay Server
-- Single-server implementation of the Skrepka protocol (v0.1)
-- Binary data (keys, signatures, blobs) uses hex encoding (bytesToHex/bytesFromHex)

-- ============================================
-- Configuration
-- ============================================

serverHost = "localhost:8080"
federationSecret = "" -- must be set to enable federation
federationScheme = "https://" -- scheme for federation HTTP calls
messageMaxAge = 86400000 -- 24 hours in milliseconds
maxPresenceTtl = 3600 -- max presence TTL from gossip in seconds

-- ============================================
-- Types
-- ============================================

type PollEvent = {id: Text, eventType: Text, payload: Text}

type GossipEvent = {eventType: Text, pubkey: Text, server: Text, timestamp: Int, ttl: Int}

-- ============================================
-- Persisted State
-- ============================================

*challenges : [{pubkey: Text, challenge: Text, expiresAt: Int}]

*sessions : [{token: Text, pubkey: Text, expiresAt: Int}]

*messages : [{messageId: Text, toKey: Text, fromKey: Text, encryptedBlob: Text, timestamp: Int, originServer: Text, receivedAt: Int}]

*presence : [{pubkey: Text, server: Text, expiresAt: Int}]

-- ============================================
-- Routes
-- ============================================

route AuthApi where
  /auth
    POST {pubkey: Text} /challenge
      -> {challenge: Text}
      = ReqChallenge
    POST {pubkey: Text, challenge: Text, signature: Text} /verify
      -> {token: Text, expiresAt: Int}
      = VerifyAuth

route ClientApi where
  GET {cursor: Int} /poll headers {authorization: Text}
    -> {events: [PollEvent], cursor: Int}
    = Poll
  /messages
    POST {to: Text, encryptedBlob: Text, timestamp: Int} / headers {authorization: Text}
      -> {status: Text, messageId: Text}
      = SendMessage
    POST {messageIds: [Text]} /ack headers {authorization: Text}
      -> {ok: Bool}
      = AckMessages
  GET /lookup/{pubkey: Text} headers {authorization: Text}
    -> {online: Bool, server: Text}
    = LookupPresence

route FederationApi where
  /federation
    POST {events: [GossipEvent], hopsRemaining: Int, fromServer: Text} /gossip headers {authorization: Text}
      -> {ok: Bool}
      = RecvGossip
    POST {toKey: Text, fromKey: Text, encryptedBlob: Text, timestamp: Int, originServer: Text, messageId: Text} /forward headers {authorization: Text}
      -> {ok: Bool}
      = ForwardMessage
    POST {messageIds: [Text], originServer: Text, fromServer: Text} /ack headers {authorization: Text}
      -> {ok: Bool}
      = FederationAck

route Api = AuthApi | ClientApi | FederationApi

-- ============================================
-- Helpers
-- ============================================

genToken = do
  a <- randomInt 2147483647
  b <- randomInt 2147483647
  c <- randomInt 2147483647
  d <- randomInt 2147483647
  yield (bytesToHex (textToBytes (show a ++ "." ++ show b ++ "." ++ show c ++ "." ++ show d)))

lookupSession = \sessions token t ->
  single (do
    s <- sessions
    where s.token == token
    where s.expiresAt > t
    yield s)

extractToken = \authHeader ->
  if take 7 authHeader == "Bearer "
    then drop 7 authHeader
    else ""

verifyFederationAuth = \authorization ->
  extractToken authorization == federationSecret && not (federationSecret == "")

clampTtl = \ttl -> if ttl < 1 then 1 else if ttl > maxPresenceTtl then maxPresenceTtl else ttl

forwardToAll = \servers path headers body ->
  if count servers == 0
    then yield False {}
    else do
      let server = fold (\acc s -> if acc == "" then s else acc) "" servers
      let rest = do s <- servers; where not (s == server); yield s
      if server == ""
        then yield False {}
        else do
          result <- httpPost (federationScheme ++ server ++ path) headers body
          restOk <- forwardToAll rest path headers body
          yield (if result.ok then True {} else restOk)

sendToEach = \servers action ->
  if count servers == 0
    then yield {}
    else do
      let server = fold (\acc s -> if acc == "" then s else acc) "" servers
      let rest = do s <- servers; where not (s == server); yield s
      if server == ""
        then yield {}
        else do
          action server
          sendToEach rest action

pruneExpired = \t -> do
  challenges <- *challenges
  set *challenges = do
    c <- challenges
    where c.expiresAt > t
    yield c
  sessions <- *sessions
  set *sessions = do
    s <- sessions
    where s.expiresAt > t
    yield s
  presence <- *presence
  set *presence = do
    p <- presence
    where p.expiresAt > t
    yield p
  messages <- *messages
  set *messages = do
    m <- messages
    where m.timestamp >= t - messageMaxAge
    yield m

-- ============================================
-- Request Handler
-- ============================================

serve = \req -> case req of

  ReqChallenge {pubkey, respond} ->
    if pubkey == ""
      then respond {challenge: ""}
      else do
        challenge <- genToken
        t <- now
        atomic do
          pruneExpired t
          challenges <- *challenges
          set *challenges = union challenges
            [{pubkey: pubkey, challenge: challenge, expiresAt: t + 60000}]
        respond {challenge: challenge}

  VerifyAuth {pubkey, challenge, signature, respond} -> do
    t <- now
    token <- genToken
    let expiry = t + 3600000
    result <- atomic do
      pruneExpired t
      challenges <- *challenges
      sessions <- *sessions
      presence <- *presence
      let matched = single (do
        c <- challenges
        where c.pubkey == pubkey
        where c.challenge == challenge
        where c.expiresAt > t
        yield c)
      let sigValid = case matched of
        Nothing {} -> False {}
        _ -> verify (bytesFromHex pubkey) (textToBytes challenge) (bytesFromHex signature)
      set *challenges = if sigValid
        then do
          c <- challenges
          where not (c.pubkey == pubkey && c.challenge == challenge)
          yield c
        else challenges
      let existingForKey = do
        s <- sessions
        where s.pubkey == pubkey
        where s.expiresAt > t
        yield s
      let oldestExpiry = fold (\acc s -> if s.expiresAt < acc then s.expiresAt else acc) expiry existingForKey
      let cappedFold = fold (\acc s ->
          if not acc.removed && s.pubkey == pubkey && s.expiresAt == oldestExpiry
            then {removed: True {}, sessions: acc.sessions}
            else {removed: acc.removed, sessions: union acc.sessions [s]}) {removed: False {}, sessions: []} sessions
      let sessionsCapped = if count existingForKey >= 10
        then cappedFold.sessions
        else sessions
      set *sessions = if sigValid
        then union sessionsCapped [{token: token, pubkey: pubkey, expiresAt: expiry}]
        else sessions
      let updatedPresence = if sigValid
        then union
          (do
            p <- presence
            where not (p.pubkey == pubkey && p.server == serverHost)
            yield p)
          [{pubkey: pubkey, server: serverHost, expiresAt: expiry}]
        else presence
      set *presence = updatedPresence
      let peerServers = if sigValid && not (federationSecret == "")
        then fold (\acc p ->
          if p.server == serverHost then acc
          else if count (do a <- acc; where a == p.server; yield a) > 0 then acc
          else union acc [p.server]) [] updatedPresence
        else []
      yield {sigValid: sigValid, peerServers: peerServers}
    if result.sigValid
      then respond {token: token, expiresAt: expiry}
      else respond {token: "", expiresAt: 0}
    if result.sigValid && not (federationSecret == "")
      then sendToEach result.peerServers (\server ->
        httpPost (federationScheme ++ server ++ "/federation/gossip")
          {authorization: "Bearer " ++ federationSecret}
          {events: [{eventType: "online", pubkey: pubkey, server: serverHost, timestamp: t, ttl: 3600}], hopsRemaining: 2, fromServer: serverHost})
      else yield {}

  Poll {cursor, authorization, respond} -> do
    t <- now
    let token = extractToken authorization
    result <- atomic do
      pruneExpired t
      sessions <- *sessions
      messages <- *messages
      let session = lookupSession sessions token t
      let evts = case session of
        Nothing {} -> []
        Just {value: s} -> do
          m <- messages
          where m.toKey == s.pubkey
          where m.timestamp >= t - messageMaxAge
          where m.receivedAt > cursor
          yield {id: m.messageId, eventType: "message",
                 payload: toJson {toKey: m.toKey, fromKey: m.fromKey,
                                  encryptedBlob: m.encryptedBlob,
                                  timestamp: m.timestamp}}
      let newCursor = case session of
        Nothing {} -> 0
        _ -> t
      yield {events: evts, cursor: newCursor}
    respond result

  SendMessage {to, encryptedBlob, timestamp, authorization, respond} -> do
    t <- now
    msgId <- genToken
    let clampedTs = if timestamp > t then t else if timestamp < t - messageMaxAge then t else timestamp
    let validInput = not (to == "") && not (encryptedBlob == "")
    result <- atomic do
      pruneExpired t
      sessions <- *sessions
      messages <- *messages
      presence <- *presence
      let session = lookupSession sessions (extractToken authorization) t
      set *messages = case session of
        Nothing {} -> messages
        Just {value: s} -> if validInput
          then union messages
            [{messageId: msgId, toKey: to, fromKey: s.pubkey,
              encryptedBlob: encryptedBlob, timestamp: clampedTs, originServer: "", receivedAt: t}]
          else messages
      let localOnline = case session of
        Nothing {} -> False {}
        Just {value: _} -> count (do
          p <- presence
          where p.pubkey == to
          where p.server == serverHost
          where p.expiresAt > t
          yield p) > 0
      let remoteServers = case session of
        Nothing {} -> []
        Just {value: _} -> do
          p <- presence
          where p.pubkey == to
          where not (p.server == serverHost)
          where p.expiresAt > t
          yield p.server
      let fromKey = case session of
        Nothing {} -> ""
        Just {value: s} -> s.pubkey
      let status = case session of
        Nothing {} -> "unauthorized"
        _ -> if not validInput then "invalid"
             else if localOnline then "delivered"
             else if count remoteServers > 0 then "federated"
             else "queued"
      yield {status: status, remoteServers: remoteServers, fromKey: fromKey}
    finalStatus <- if result.status == "federated" && not (federationSecret == "")
      then do
        anyOk <- forwardToAll result.remoteServers "/federation/forward"
          {authorization: "Bearer " ++ federationSecret}
          {toKey: to, fromKey: result.fromKey, encryptedBlob: encryptedBlob,
           timestamp: clampedTs, originServer: serverHost, messageId: msgId}
        yield (if anyOk then "federated" else "queued")
      else yield result.status
    respond {status: finalStatus, messageId: (if result.status == "unauthorized" then "" else if result.status == "invalid" then "" else msgId)}

  AckMessages {messageIds, authorization, respond} -> do
    t <- now
    let dedupedIds = fold (\acc mid ->
      if count (do a <- acc; where a == mid; yield a) > 0 then acc
      else union acc [mid]) [] messageIds
    result <- atomic do
      pruneExpired t
      sessions <- *sessions
      messages <- *messages
      let session = lookupSession sessions (extractToken authorization) t
      let ackedMessages = case session of
        Nothing {} -> []
        Just {value: s} -> do
          m <- messages
          mid <- dedupedIds
          where m.messageId == mid
          where m.toKey == s.pubkey
          yield m
      set *messages = case session of
        Nothing {} -> messages
        Just {value: _} -> diff messages ackedMessages
      let ok = case session of
        Nothing {} -> False {}
        _ -> True {}
      yield {ok: ok, ackedMessages: ackedMessages}
    let fedMessages = do
      m <- result.ackedMessages
      where not (m.originServer == "")
      where not (m.originServer == serverHost)
      yield m
    let originServers = fold (\acc m ->
      if count (do a <- acc; where a == m.originServer; yield a) > 0 then acc
      else union acc [m.originServer]) [] fedMessages
    if result.ok && count originServers > 0 && not (federationSecret == "")
      then sendToEach originServers (\server ->
        httpPost (federationScheme ++ server ++ "/federation/ack")
          {authorization: "Bearer " ++ federationSecret}
          {messageIds: (do m <- fedMessages; where m.originServer == server; yield m.messageId),
           originServer: server, fromServer: serverHost})
      else yield {}
    respond {ok: result.ok}

  LookupPresence {pubkey, authorization, respond} -> do
    t <- now
    result <- atomic do
      pruneExpired t
      sessions <- *sessions
      presence <- *presence
      let session = lookupSession sessions (extractToken authorization) t
      let valid = do
        p <- presence
        where p.pubkey == pubkey
        where p.expiresAt > t
        yield p
      let foundServer = case session of
        Nothing {} -> ""
        Just {value: _} ->
          if count valid == 0
            then ""
            else case single (do
              p <- valid
              where p.server == serverHost
              yield p) of
              Just {value: p} -> p.server
              Nothing {} -> fold (\_ p -> p.server) "" valid
      let authorized = case session of
        Nothing {} -> False {}
        _ -> True {}
      yield {foundServer: foundServer, authorized: authorized}
    if result.authorized
      then respond {online: not (result.foundServer == ""), server: result.foundServer}
      else respond {online: False {}, server: "unauthorized"}

  RecvGossip {events, hopsRemaining, fromServer, authorization, respond} -> do
    t <- now
    let authorized = verifyFederationAuth authorization
    let cappedHops = if hopsRemaining > 5 then 5 else hopsRemaining
    let validRequest = authorized && cappedHops >= 1
    let onlineEntries = do
      e <- events
      where e.eventType == "online"
      where not (e.server == "")
      where not (e.pubkey == "")
      where not (e.server == serverHost)
      where e.timestamp <= t
      where e.timestamp + (clampTtl e.ttl) * 1000 > t
      yield {pubkey: e.pubkey, server: e.server, expiresAt: e.timestamp + (clampTtl e.ttl) * 1000, timestamp: e.timestamp}
    let offlineEntries = do
      e <- events
      where e.eventType == "offline"
      where not (e.server == "")
      where not (e.pubkey == "")
      where not (e.server == serverHost)
      where e.timestamp <= t
      where t - e.timestamp < (clampTtl e.ttl) * 1000
      yield {pubkey: e.pubkey, server: e.server, timestamp: e.timestamp}
    let filteredOnline = do
      e <- onlineEntries
      where count (do
        o <- offlineEntries
        where o.pubkey == e.pubkey
        where o.server == e.server
        where o.timestamp >= e.timestamp
        yield o) == 0
      yield {pubkey: e.pubkey, server: e.server, expiresAt: e.expiresAt}
    let affectedKeys = union
      (do
        e <- onlineEntries
        yield {pubkey: e.pubkey, server: e.server})
      (do
        e <- offlineEntries
        yield {pubkey: e.pubkey, server: e.server})
    let deduped = fold (\acc e ->
        if count (do
            a <- acc
            where a.pubkey == e.pubkey
            where a.server == e.server
            where a.expiresAt >= e.expiresAt
            yield a) > 0
          then acc
          else union
            (do
              a <- acc
              where not (a.pubkey == e.pubkey && a.server == e.server)
              yield a)
            [e]) [] filteredOnline
    atomic do
      pruneExpired t
      presence <- *presence
      let cleaned = diff presence (do
        p <- presence
        k <- affectedKeys
        where p.pubkey == k.pubkey
        where p.server == k.server
        yield p)
      set *presence = if validRequest
        then union cleaned deduped
        else presence
    let validEvents = do
      e <- events
      where not (e.server == "")
      where not (e.pubkey == "")
      where e.timestamp <= t
      where e.timestamp + (clampTtl e.ttl) * 1000 > t
      yield e
    if validRequest && cappedHops > 1 && not (federationSecret == "") && count validEvents > 0
      then do
        peerServers <- atomic do
          presence <- *presence
          yield (fold (\acc p ->
            if p.server == serverHost then acc
            else if count (do a <- acc; where a == p.server; yield a) > 0 then acc
            else union acc [p.server]) [] presence)
        let senderServers = fold (\acc e ->
          if count (do a <- acc; where a == e.server; yield a) > 0 then acc
          else union acc [e.server]) (if fromServer == "" then [] else [fromServer]) events
        let targets = do
          s <- peerServers
          where count (do ss <- senderServers; where ss == s; yield ss) == 0
          yield s
        sendToEach targets (\server ->
          httpPost (federationScheme ++ server ++ "/federation/gossip")
            {authorization: "Bearer " ++ federationSecret}
            {events: validEvents, hopsRemaining: cappedHops - 1, fromServer: serverHost})
      else yield {}
    respond {ok: validRequest}

  ForwardMessage {toKey, fromKey, encryptedBlob, timestamp, originServer, messageId, authorization, respond} -> do
    t <- now
    let authorized = verifyFederationAuth authorization
    let validOrigin = not (originServer == "") && not (originServer == serverHost)
    let validFields = not (toKey == "") && not (fromKey == "") && not (encryptedBlob == "") && not (messageId == "")
    let clampedTs = if timestamp > t then t else if timestamp < t - messageMaxAge then t else timestamp
    atomic do
      pruneExpired t
      messages <- *messages
      let exists = count (do
        m <- messages
        where m.messageId == messageId
        where m.timestamp >= t - messageMaxAge
        yield m) > 0
      let shouldInsert = authorized && validOrigin && validFields && not exists
      set *messages = if shouldInsert
        then union messages
          [{messageId: messageId, toKey: toKey, fromKey: fromKey,
            encryptedBlob: encryptedBlob, timestamp: clampedTs, originServer: originServer, receivedAt: t}]
        else messages
    respond {ok: authorized && validOrigin && validFields}

  FederationAck {messageIds, originServer, fromServer, authorization, respond} -> do
    t <- now
    let authorized = verifyFederationAuth authorization
    let validAck = authorized && originServer == serverHost && not (fromServer == "") && not (fromServer == serverHost)
    atomic do
      pruneExpired t
      messages <- *messages
      presence <- *presence
      set *messages = if validAck
        then diff messages (do
          m <- messages
          mid <- messageIds
          where m.messageId == mid
          where count (do
            p <- presence
            where p.pubkey == m.toKey
            where p.server == serverHost
            where p.expiresAt > t
            yield p) == 0
          yield m)
        else messages
    respond {ok: validAck}

-- ============================================
-- Entry Point
-- ============================================

main = do
  println "Skrepka relay server starting on port 8080..."
  listen 8080 serve