# Copyright 2019 The Cobalt Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# pylint: disable=redefined-outer-name,redefined-builtin,broad-except,unused-argument
'''Collectd module for Cobalt'''

import json
import time
import argparse
import websocket

try:
  import collectd
except ImportError:
  collectd = None

conf_ = {'host': 'localhost', 'port': 9222}


def config(conf):
  print(f'config called {conf!r}')
  conf_['collectd_conf'] = conf
  collectd.info(f'config called {conf!r}')


def reconnect():
  ws = websocket.create_connection(f"ws://{conf_['host']}:{conf_['port']}")
  ws.settimeout(3)
  setattr(ws, 'message_id', 1)
  conf_['ws'] = ws


def init():
  conf = conf_['collectd_conf']
  for child in conf.children:
    collectd.info(f'conf.child key {child.key} values {child.values}')
    if child.key.lower() == 'address':
      host, port = child.values[0].split(':')
      conf_['host'] = host
      conf_['port'] = int(port)

  reconnect()


def wait_result(ws, result_id, timeout):
  start_time = time.time()
  messages = []
  matching_result = None
  while True:
    now = time.time()
    if now - start_time > timeout:
      break
    try:
      message = ws.recv()
      parsed_message = json.loads(message)
      messages.append(parsed_message)
      if 'result' in parsed_message and parsed_message['id'] == result_id:
        matching_result = parsed_message
        break
    except Exception:
      break
  return (matching_result, messages)


def exec_js(ws, expr):
  ws.message_id += 1
  call_obj = {
      'id': ws.message_id,
      'method': 'Runtime.evaluate',
      'params': {
          'expression': expr,
          'returnByValue': True
      }
  }
  ws.send(json.dumps(call_obj))
  result, _ = wait_result(ws, ws.message_id, 10)
  return result['result']['result']


def report_value(plugin, type, instance, value, plugin_instance=None):
  vl = collectd.Values(plugin=plugin, type=type, type_instance=instance)
  if plugin_instance:
    vl.plugin_instance = plugin_instance
  vl.dispatch(values=[value])


def report_cobalt_stat(key, collectd_type):
  val = exec_js(conf_['ws'], 'h5vcc.cVal.getValue("' + key + '")')
  try:
    value = int(val['value'])
    # collectd.warning('key {} value: {}'.format(key, value))
    report_value('cobalt', collectd_type, key, value)
  except TypeError:
    collectd.warning(f'Failed to collect: {key} {val}')


# Tuple of a Cobalt Cval, and data type
# Collectd frontends like CGP group values to graphs by the reported data type.
# The types here are simple gauges grouped by their rough estimated orders of
# magnitude, so that all plots are still somewhat discernible.
tracked_stats = [('Cobalt.Lifetime', 'lifetime'),
                 ('Count.DOM.ActiveJavaScriptEvents', 'gauge'),
                 ('Count.DOM.Attrs', 'gauge_10k'),
                 ('Count.DOM.EventListeners', 'gauge_10k'),
                 ('Count.DOM.HtmlCollections', 'gauge_10k'),
                 ('Count.DOM.NodeLists', 'gauge_10k'),
                 ('Count.DOM.NodeMaps', 'gauge_10k'),
                 ('Count.DOM.Nodes', 'gauge_10k'),
                 ('Count.DOM.StringMaps', 'gauge_100'),
                 ('Count.DOM.TokenLists', 'gauge_10k'),
                 ('Count.MainWebModule.DOM.HtmlElement', 'gauge_10k'),
                 ('Count.MainWebModule.DOM.HtmlElement.Document', 'gauge_10k'),
                 ('Count.MainWebModule.DOM.HtmlScriptElement.Execute', 'gauge'),
                 ('Count.Renderer.Rasterize.NewRenderTree', 'gauge_100k'),
                 ('Count.XHR', 'gauge_100'), ('Memory.CPU.Free', 'gauge_1G'),
                 ('Memory.CPU.Used', 'gauge_100M'),
                 ('Memory.DebugConsoleWebModule.DOM.HtmlScriptElement.Execute',
                  'gauge_100k'),
                 ('Memory.Font.LocalTypefaceCache.Capacity', 'gauge_10M'),
                 ('Memory.Font.LocalTypefaceCache.Size', 'gauge_10M'),
                 ('Memory.JS', 'gauge_10M'),
                 ('Memory.MainWebModule.DOM.HtmlScriptElement.Execute',
                  'gauge_1M'), ('Memory.XHR', 'gauge_1M'),
                 ('Memory.MainWebModule.ImageCache.Size', 'gauge_10M'),
                 ('Renderer.SubmissionQueueSize', 'gauge')]


def cobalt_read():
  for name, type in tracked_stats:
    report_cobalt_stat(name, type)


def read(*args, **kwargs):
  try:
    cobalt_read()
  except Exception:
    reconnect()


if collectd:
  collectd.register_config(config)
  collectd.register_init(init)
  collectd.register_read(read)

# Debugcode to verify plugin connection
if __name__ == '__main__':

  class Bunch:

    def __init__(self, **kwds):
      self.__dict__.update(kwds)

  def debugprint(*args, **kwargs):
    print(f'{args!r} {kwargs!r}')

  def debugvalues(**kwargs):
    debugprint(**kwargs)
    return Bunch(dispatch=debugprint, **kwargs)

  collectd = Bunch(info=debugprint, warning=debugprint, Values=debugvalues)
  parser = argparse.ArgumentParser()
  parser.add_argument('--host', default='localhost')
  parser.add_argument('--port', type=int, default=9222)
  args = parser.parse_args()
  conf = Bunch(
      children=[Bunch(key='address', values=[f'{args.host}:{args.port}'])])
  conf_['collectd_conf'] = conf
  init()
  cobalt_read()