import os
import sys
import json
import aiohttp
import logging

from collections import namedtuple
from functools import reduce
from aiohttp import web

PORT = os.getenv("PORT", 9115)

UP_HELP = "# HELP state_up State Up/Down status."
UP_TYPE = "# TYPE state_up gauge"
UP_FORMAT = "state_up %d"

CONNECTED_HELP = "# HELP state_connected Connected status."
CONNECTED_TYPE = "# TYPE state_connected gauge"
CONNECTED_FORMAT = "state_connected{display_id=\"%d\"} %d"

LAST_DURATION_HELP = "# HELP state_last_duration Last duration."
LAST_DURATION_TYPE = "# TYPE state_last_duration gauge"
LAST_DURATION_FORMAT = "state_last_duration{display_id=\"%d\"} %d"

LAST_FORECASTS_COUNT_HELP = "# HELP state_last_forecasts_count Last forecasts count."
LAST_FORECASTS_COUNT_TYPE = "# TYPE state_last_forecasts_count gauge"
LAST_FORECASTS_COUNT_FORMAT = "state_last_forecasts_count{display_id=\"%d\"} %d"

Metrics = namedtuple("Metrics", "connected duration forecasts")

def reducer(metrics, state):
  display_id = int(state["DisplayId"])

  forecasts = -1
  if state["Connected"]:
    forecasts = int(state["LastForecastsCount"])

  metrics.connected.append(CONNECTED_FORMAT % (display_id, int(state["Connected"])))
  metrics.duration.append(LAST_DURATION_FORMAT % (display_id, int(state["LastDuration"])))
  metrics.forecasts.append(LAST_FORECASTS_COUNT_FORMAT % (display_id, forecasts))

  return metrics

def format_metrics(up, metrics):
  if up:
    return "\n".join([
      UP_HELP,
      UP_TYPE,
      UP_FORMAT % int(up),
      CONNECTED_HELP,
      CONNECTED_TYPE,
      "\n".join(metrics.connected),
      LAST_DURATION_HELP,
      LAST_DURATION_TYPE,
      "\n".join(metrics.duration),
      LAST_FORECASTS_COUNT_HELP,
      LAST_FORECASTS_COUNT_TYPE,
      "\n".join(metrics.forecasts),
    ])

  return "\n".join([UP_HELP, UP_TYPE, UP_FORMAT % int(up)])


async def persistent_session(app):
  app["PERSISTENT_SESSION"] = aiohttp.ClientSession()
  yield
  await app["PERSISTENT_SESSION"].close()

async def fetch(session, url):
  async with session.get(url) as response:
    data = await response.read()
    return json.loads(data)

async def handle(request):
  try:
    target = request.query["target"]
    session = request.app["PERSISTENT_SESSION"]
    data = await fetch(session, target)
    states = data["GetDisplaysStatesResult"]["DisplayState"]
    metrics = reduce(reducer, states, Metrics([], [], []))

    return web.Response(text=format_metrics(True, metrics))
  except:
    return web.Response(text=format_metrics(False, None))

logging.basicConfig(stream=sys.stdout)

app = web.Application()
app.router.add_get('/probe', handle)
app.cleanup_ctx.append(persistent_session)

if __name__ == '__main__':
  web.run_app(app, port=PORT)