# Copyright 2018 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import contextlib
import logging
import os
import shutil
import subprocess
import sys
import tempfile

from devil import devil_env
from devil.android import device_signal, device_errors
from devil.android.sdk import version_codes
from pylib import constants


def _ProcessType(proc):
    _, _, suffix = proc.name.partition(':')
    if not suffix:
        return 'browser'
    if suffix.startswith('sandboxed_process'):
        return 'renderer'
    if suffix.startswith('privileged_process'):
        return 'gpu'
    return None


def _GetSpecifiedPID(device, package_name, process_specifier):
    if process_specifier is None:
        return None

    # Check for numeric PID
    try:
        pid = int(process_specifier)
        return pid
    except ValueError:
        pass

    # Check for exact process name; can be any of these formats:
    #   <package>:<process name>, i.e. 'org.chromium.chrome:sandboxed_process0'
    #   :<process name>, i.e. ':sandboxed_process0'
    #   <process name>, i.e. 'sandboxed_process0'
    full_process_name = process_specifier
    if process_specifier.startswith(':'):
        full_process_name = package_name + process_specifier
    elif ':' not in process_specifier:
        full_process_name = '%s:%s' % (package_name, process_specifier)
    matching_processes = device.ListProcesses(full_process_name)
    if len(matching_processes) == 1:
        return matching_processes[0].pid
    if len(matching_processes) > 1:
        raise RuntimeError(
            'Found %d processes with name "%s".'
            % (len(matching_processes), process_specifier)
        )

    # Check for process type (i.e. 'renderer')
    package_processes = device.ListProcesses(package_name)
    matching_processes = [
        p for p in package_processes if (_ProcessType(p) == process_specifier)
    ]
    if process_specifier == 'renderer' and len(matching_processes) > 1:
        raise RuntimeError(
            'Found %d renderer processes; please re-run with only '
            'one open tab.' % len(matching_processes)
        )
    if len(matching_processes) != 1:
        raise RuntimeError(
            'Found %d processes of type "%s".'
            % (len(matching_processes), process_specifier)
        )
    return matching_processes[0].pid


def _ThreadsForProcess(device, pid):
    # The thread list output format for 'ps' is the same regardless of version.
    # Here's the column headers, and a sample line for a thread belonging to
    # pid 12345 (note that the last few columns are not aligned with headers):
    #
    # USER        PID   TID  PPID     VSZ    RSS WCHAN            ADDR S CMD
    # u0_i101   12345 24680   567 1357902  97531 futex_wait_queue_me e85acd9c S \
    #     CrRendererMain
    if device.build_version_sdk >= version_codes.OREO:
        pid_regex = (
            r'^[[:graph:]]\{1,\}[[:blank:]]\{1,\}%d[[:blank:]]\{1,\}' % pid
        )
        ps_cmd = "ps -T -e | grep '%s'" % pid_regex
        ps_output_lines = device.RunShellCommand(
            ps_cmd, shell=True, check_return=True
        )
    else:
        ps_cmd = ['ps', '-p', str(pid), '-t']
        ps_output_lines = device.RunShellCommand(ps_cmd, check_return=True)
    result = []
    for l in ps_output_lines:
        fields = l.split()
        # fields[2] is tid, fields[-1] is thread name. Output may include an entry
        # for the process itself with tid=pid; omit that one.
        if fields[2] == str(pid):
            continue
        result.append((int(fields[2]), fields[-1]))
    return result


def _ThreadType(thread_name):
    if not thread_name:
        return 'unknown'
    if thread_name.startswith('Chrome_ChildIO') or thread_name.startswith(
        'Chrome_IO'
    ):
        return 'io'
    if thread_name.startswith('Compositor'):
        return 'compositor'
    if (
        thread_name.startswith('ChildProcessMai')
        or thread_name.startswith('CrGpuMain')
        or thread_name.startswith('CrRendererMain')
    ):
        return 'main'
    if thread_name.startswith('RenderThread'):
        return 'render'
    raise ValueError('got no matching thread_name')


def _GetSpecifiedTID(device, pid, thread_specifier):
    if thread_specifier is None:
        return None

    # Check for numeric TID
    try:
        tid = int(thread_specifier)
        return tid
    except ValueError:
        pass

    # Check for thread type
    if pid is not None:
        matching_threads = [
            t
            for t in _ThreadsForProcess(device, pid)
            if (_ThreadType(t[1]) == thread_specifier)
        ]
        if len(matching_threads) != 1:
            raise RuntimeError(
                'Found %d threads of type "%s".'
                % (len(matching_threads), thread_specifier)
            )
        return matching_threads[0][0]

    return None


def PrepareDevice(device):
    if device.build_version_sdk < version_codes.NOUGAT:
        raise RuntimeError(
            'Simpleperf profiling is only supported on Android N and later.'
        )

    # Necessary for profiling
    # https://android-review.googlesource.com/c/platform/system/sepolicy/+/234400
    device.SetProp('security.perf_harden', '0')


def InstallSimpleperf(device, package_name):
    package_arch = device.GetPackageArchitecture(package_name) or 'armeabi-v7a'
    host_simpleperf_path = devil_env.config.LocalPath(
        'simpleperf', package_arch
    )
    if not host_simpleperf_path:
        raise Exception('Could not get path to simpleperf executable on host.')
    device_simpleperf_path = '/'.join(
        ('/data/local/tmp/profilers', package_arch, 'simpleperf')
    )
    device.PushChangedFiles([(host_simpleperf_path, device_simpleperf_path)])
    return device_simpleperf_path


@contextlib.contextmanager
def RunSimpleperf(
    device,
    device_simpleperf_path,
    package_name,
    process_specifier,
    thread_specifier,
    events,
    profiler_args,
    host_out_path,
):
    pid = _GetSpecifiedPID(device, package_name, process_specifier)
    tid = _GetSpecifiedTID(device, pid, thread_specifier)
    if pid is None and tid is None:
        raise RuntimeError(
            'Could not find specified process/thread running on '
            'device. Make sure the apk is already running before '
            'attempting to profile.'
        )
    profiler_args = list(profiler_args)
    if profiler_args and profiler_args[0] == 'record':
        profiler_args.pop(0)
    profiler_args.extend(('-e', events))
    if '--call-graph' not in profiler_args and '-g' not in profiler_args:
        profiler_args.append('-g')
    if '-f' not in profiler_args:
        profiler_args.extend(('-f', '1000'))

    device_out_path = '/data/local/tmp/perf.data'
    should_remove_device_out_path = True
    if '-o' in profiler_args:
        device_out_path = profiler_args[profiler_args.index('-o') + 1]
        should_remove_device_out_path = False
    else:
        profiler_args.extend(('-o', device_out_path))

    # Remove the default output to avoid confusion if simpleperf opts not
    # to update the file.
    file_exists = True
    try:
        device.adb.Shell('readlink -e ' + device_out_path)
    except device_errors.AdbCommandFailedError:
        file_exists = False
    if file_exists:
        logging.warning(
            '%s output file already exists on device', device_out_path
        )
        if not should_remove_device_out_path:
            raise RuntimeError(
                'Specified output file \'{}\' already exists, not '
                'continuing'.format(device_out_path)
            )
    device.adb.Shell('rm -f ' + device_out_path)

    if tid:
        profiler_args.extend(('-t', str(tid)))
    else:
        profiler_args.extend(('-p', str(pid)))

    adb_shell_simpleperf_process = device.adb.StartShell(
        [device_simpleperf_path, 'record'] + profiler_args
    )

    completed = False
    try:
        yield
        completed = True

    finally:
        device.KillAll(
            'simpleperf', signum=device_signal.SIGINT, blocking=True, quiet=True
        )
        if completed:
            adb_shell_simpleperf_process.wait()
            ret = adb_shell_simpleperf_process.returncode
            if ret == 0:
                # Successfully gathered a profile
                device.PullFile(device_out_path, host_out_path)
            else:
                logging.warning(
                    'simpleperf exited unusually, expected exit 0, got %d', ret
                )
                stdout, stderr = adb_shell_simpleperf_process.communicate()
                logging.info('stdout: \'%s\', stderr: \'%s\'', stdout, stderr)
                raise RuntimeError(
                    'simpleperf exited with unexpected code {} '
                    '(run with -vv for full stdout/stderr)'.format(ret)
                )


def ConvertSimpleperfToPprof(
    simpleperf_out_path, build_directory, pprof_out_path
):
    # The simpleperf scripts require the unstripped libs to be installed in the
    # same directory structure as the libs on the device. Much of the logic here
    # is just figuring out and creating the necessary directory structure, and
    # symlinking the unstripped shared libs.

    # Get the set of libs that we can symbolize
    unstripped_lib_dir = os.path.join(build_directory, 'lib.unstripped')
    unstripped_libs = set(
        f for f in os.listdir(unstripped_lib_dir) if f.endswith('.so')
    )

    # report.py will show the directory structure above the shared libs;
    # that is the directory structure we need to recreate on the host.
    script_dir = devil_env.config.LocalPath('simpleperf_scripts')
    report_path = os.path.join(script_dir, 'report.py')
    report_cmd = [sys.executable, report_path, '-i', simpleperf_out_path]
    device_lib_path = None
    output = subprocess.check_output(report_cmd, stderr=subprocess.STDOUT)
    if isinstance(output, bytes):
        output = output.decode()
    for line in output.splitlines():
        fields = line.split()
        if len(fields) < 5:
            continue
        shlib_path = fields[4]
        shlib_dirname, shlib_basename = shlib_path.rpartition('/')[::2]
        if shlib_basename in unstripped_libs:
            device_lib_path = shlib_dirname
            break
    if not device_lib_path:
        raise RuntimeError(
            'No chrome-related symbols in profiling data in %s. '
            'Either the process was idle for the entire profiling '
            'period, or something went very wrong (and you should '
            'file a bug at crbug.com/new with component '
            'Speed>Tracing, and assign it to szager@chromium.org).'
            % simpleperf_out_path
        )

    # Recreate the directory structure locally, and symlink unstripped libs.
    processing_dir = tempfile.mkdtemp()
    try:
        processing_lib_dir = os.path.join(
            processing_dir, 'binary_cache', device_lib_path.lstrip('/')
        )
        os.makedirs(processing_lib_dir)
        for lib in unstripped_libs:
            unstripped_lib_path = os.path.join(unstripped_lib_dir, lib)
            processing_lib_path = os.path.join(processing_lib_dir, lib)
            os.symlink(unstripped_lib_path, processing_lib_path)

        # Run the script to annotate symbols and convert from simpleperf format to
        # pprof format.
        pprof_converter_script = os.path.join(
            script_dir, 'pprof_proto_generator.py'
        )
        pprof_converter_cmd = [
            sys.executable,
            pprof_converter_script,
            '-i',
            simpleperf_out_path,
            '-o',
            os.path.abspath(pprof_out_path),
            '--ndk_path',
            constants.ANDROID_NDK_ROOT,
        ]
        subprocess.check_output(
            pprof_converter_cmd, stderr=subprocess.STDOUT, cwd=processing_dir
        )
    finally:
        shutil.rmtree(processing_dir, ignore_errors=True)