#!/usr/bin/env vpython3

# Copyright 2024 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""The module to create and manage measures using in the process."""

import functools
import json
import os
import sys

from google.protobuf import any_pb2
from google.protobuf.json_format import MessageToDict

# Add to sys.path so that this module can be imported by other modules that
# have different path setup, e.g. android test runner, and ios test runner.
sys.path.append(os.path.abspath(os.path.dirname(__file__)))
from average import Average
from count import Count
from data_points import DataPoints
from measure import Measure
from metric import Metric
from time_consumption import TimeConsumption

# This is used as the key when being uploaded to ResultDB via result_sink
# and shouldn't be changed
TEST_SCRIPT_METRICS_KEY = 'test_script_metrics'

# The file name is used as the key when being loaded into the ResultDB and
# shouldn't be changed.
TEST_SCRIPT_METRICS_JSONPB_FILENAME = f'{TEST_SCRIPT_METRICS_KEY}.jsonpb'

_metric = Metric()


def _create_name(*name_pieces: str) -> str:
    if len(name_pieces) == 0:
        raise ValueError('Need at least one name piece.')
    return '/'.join(list(name_pieces))


def _register(m: Measure) -> Measure:
    _metric.register(m)
    return m


def average(*name_pieces: str) -> Average:
    return _register(Average(_create_name(*name_pieces)))


def count(*name_pieces: str) -> Count:
    return _register(Count(_create_name(*name_pieces)))


def data_points(*name_pieces: str) -> DataPoints:
    return _register(DataPoints(_create_name(*name_pieces)))


def time_consumption(*name_pieces: str) -> TimeConsumption:
    return _register(TimeConsumption(_create_name(*name_pieces)))


def timed_func(*name_pieces: str):
    """time_consumption() as a @decorator."""

    def decorator(func):

        @functools.wraps(func)
        def wrapped(*args, **kwargs):
            with time_consumption(*name_pieces):
                func(*args, **kwargs)

        return wrapped

    return decorator


def tag(*args: str) -> None:
    """Adds a tag to the Metric to tag the final results; see Metric for details."""
    _metric.tag(*args)


def clear() -> None:
    """Clears all the registered Measures."""
    _metric.clear()


def size() -> int:
    """Gets the current size of registered Measures."""
    return _metric.size()


def to_dict() -> dict:
    """Converts all the registered Measures to a dict.

    The records are wrapped in protobuf Any message before exported as dict
    so that an additional key "@type" is included.
    """
    any_msg = any_pb2.Any()
    any_msg.Pack(_metric.dump())
    return MessageToDict(any_msg, preserving_proto_field_name=True)


def to_json() -> str:
    """Converts all the registered Measures to a json str."""
    return json.dumps(to_dict(), sort_keys=True, indent=2)


# TODO(crbug.com/343242386): May need to implement a lock and reset logic to
# clear in-memory data and lock the instance to block further operations and
# avoid accidentally accumulating data which won't be published at all.
def dump(dir_path: str) -> None:
    """Dumps the metric data into test_script_metrics.jsonpb in the |path|."""
    os.makedirs(dir_path, exist_ok=True)
    with open(
        os.path.join(dir_path, TEST_SCRIPT_METRICS_JSONPB_FILENAME),
        'w',
        encoding='utf-8',
    ) as wf:
        wf.write(to_json())