#!/usr/bin/env python
# Copyright 2014 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from skypy.skyserver import SkyServer
import argparse
import json
import logging
import os
import pipes
import re
import requests
import signal
import skypy.paths
import StringIO
import subprocess
import sys
import time
import urlparse

SUPPORTED_MIME_TYPES = [
    'text/html',
    'text/sky',
    'text/plain',
]

DEFAULT_SKY_COMMAND_PORT = 7777
GDB_PORT = 8888
SKY_SERVER_PORT = 9999
DEFAULT_URL = "https://raw.githubusercontent.com/domokit/mojo/master/sky/examples/home.sky"

ANDROID_PACKAGE = "org.chromium.mojo.shell"
ANDROID_ACTIVITY = "%s/.MojoShellActivity" % ANDROID_PACKAGE
ANDROID_APK_NAME = 'MojoShell.apk'

PID_FILE_PATH = "/tmp/skydb.pids"
CACHE_LINKS_PATH = '/tmp/mojo_cache_links'
SYSTEM_LIBS_ROOT_PATH = '/tmp/device_libs'

SRC_ROOT = skypy.paths.Paths('ignored').src_root
ADB_PATH = os.path.join(SRC_ROOT,
    'third_party/android_tools/sdk/platform-tools/adb')


# FIXME: Move this into mopy.config
def gn_args_from_build_dir(build_dir):
    gn_cmd = [
        'gn', 'args',
        build_dir,
        '--list', '--short'
    ]
    config = {}
    for line in subprocess.check_output(gn_cmd).strip().split('\n'):
        # FIXME: This doesn't handle = in values.
        key, value = line.split(' = ')
        config[key] = value
    return config


class SkyDebugger(object):
    def __init__(self):
        self.pids = {}
        self.paths = None

    def _server_root_for_url(self, url_or_path):
        path = os.path.abspath(url_or_path)
        if os.path.commonprefix([path, SRC_ROOT]) == SRC_ROOT:
            server_root = SRC_ROOT
        else:
            server_root = os.path.dirname(path)
            logging.warn(
                '%s is outside of mojo root, using %s as server root' %
                (path, server_root))
        return server_root

    def _in_chromoting(self):
        return os.environ.get('CHROME_REMOTE_DESKTOP_SESSION', False)

    def _wrap_for_android(self, shell_args):
        # am shell --esa: (someone shoot me now)
        #  [--esa <EXTRA_KEY> <EXTRA_STRING_VALUE>[,<EXTRA_STRING_VALUE...]]
        #  (to embed a comma into a string escape it using "\,")
        escaped_args = map(lambda arg: arg.replace(',', '\\,'), shell_args)
        return [
            ADB_PATH, 'shell',
            'am', 'start',
            '-W',
            '-S',
            '-a', 'android.intent.action.VIEW',
            '-n', ANDROID_ACTIVITY,
            # FIXME: This quoting is very error-prone.  Perhaps we should read
            # our args from a file instead?
            '--esa', 'parameters', ','.join(escaped_args),
        ]

    def _build_mojo_shell_command(self, args, is_android):
        content_handlers = ['%s,%s' % (mime_type, 'mojo:sky_viewer')
            for mime_type in SUPPORTED_MIME_TYPES]

        remote_command_port = self.pids.get('remote_sky_command_port', self.pids['sky_command_port'])
        remote_server_port = self.pids.get('remote_sky_server_port', self.pids['sky_server_port'])

        shell_args = [
            '--v=1',
            '--content-handlers=%s' % ','.join(content_handlers),
            '--url-mappings=mojo:window_manager=mojo:kiosk_wm',
            '--args-for=mojo:sky_debugger %d' % remote_command_port,
            'mojo:sky_debugger',
        ]

        if args.url_or_path:
            shell_args.append(
                '--args-for=mojo:window_manager %s' % self._url_from_args(args))

        # Map all mojo: urls to http: urls using the --origin command.
        build_dir_url = SkyServer.url_for_path(
            remote_server_port,
            self.pids['sky_server_root'],
            self.pids['build_dir'])

        # TODO(eseidel): We should do this on linux, but we need to fix
        # mojo http loading to be faster first.
        if is_android:
            shell_args += ['--origin=%s' % build_dir_url]

        # Desktop-only work-around for mojo crashing under chromoting.
        if not is_android and args.use_osmesa:
            shell_args.append(
                '--args-for=mojo:native_viewport_service --use-osmesa')

        if is_android and args.gdb:
            shell_args.append('--wait-for-debugger')

        if 'remote_sky_server_port' in self.pids:
            shell_command = self._wrap_for_android(shell_args)
        else:
            shell_command = [self.paths.mojo_shell_path] + shell_args

        return shell_command

    def sky_server_for_args(self, args):
        # FIXME: This is a hack.  sky_server should just take a build_dir
        # not a magical "configuration" name.
        configuration = os.path.basename(os.path.normpath(self.paths.build_dir))
        server_root = self._server_root_for_url(args.url_or_path)
        sky_server = SkyServer(SKY_SERVER_PORT, configuration, server_root)
        return sky_server

    def _create_paths_for_build_dir(self, build_dir):
        # skypy.paths.Paths takes a root-relative build_dir argument. :(
        abs_build_dir = os.path.abspath(build_dir)
        root_relative_build_dir = os.path.relpath(abs_build_dir, SRC_ROOT)
        return skypy.paths.Paths(root_relative_build_dir)

    def _find_remote_pid_for_package(self, package):
        ps_output = subprocess.check_output([ADB_PATH, 'shell', 'ps'])
        for line in ps_output.split('\n'):
            fields = line.split()
            if fields and fields[-1] == package:
                return fields[1]
        return None

    def _find_install_location_for_package(self, package):
        pm_command = [ADB_PATH, 'shell', 'pm', 'path', package]
        pm_output = subprocess.check_output(pm_command)
        # e.g. package:/data/app/org.chromium.mojo.shell-1/base.apk
        return pm_output.split(':')[-1]

    def start_command(self, args):
        # FIXME: Lame that we use self for a command-specific variable.
        self.paths = self._create_paths_for_build_dir(args.build_dir)
        self.stop_command(None) # Quit any existing process.

        if not os.path.exists(self.paths.mojo_shell_path):
            print "mojo_shell not found in build_dir '%s'" % args.build_dir
            print "Are you sure you sure that's a valid build_dir location?"
            print "See skydb start --help for more info"
            sys.exit(2)

        # FIXME: This is probably not the right way to compute is_android
        # from the build directory?
        gn_args = gn_args_from_build_dir(self.paths.build_dir)
        is_android = 'android_sdk_version' in gn_args

        if is_android and args.gdb and not 'is_debug' in gn_args:
            # FIXME: We don't include gdbserver in the release APK...
            print "Cannot debug Release builds on Android"
            sys.exit(2)

        sky_server = self.sky_server_for_args(args)
        self.pids['sky_server_pid'] = sky_server.start()
        self.pids['sky_server_port'] = sky_server.port
        self.pids['sky_server_root'] = sky_server.root

        self.pids['build_dir'] = self.paths.build_dir
        self.pids['sky_command_port'] = args.command_port

        if is_android:
            # TODO(eseidel): This should move into a helper method and handle
            # failures with nice messages explaining how to get root.
            subprocess.check_call([ADB_PATH, 'root'])

            # We could make installing conditional on an argument.
            apk_path = os.path.join(self.paths.build_dir, 'apks',
                ANDROID_APK_NAME)
            subprocess.check_call([ADB_PATH, 'install', '-r', apk_path])

            port_string = 'tcp:%s' % sky_server.port
            subprocess.check_call([
                ADB_PATH, 'reverse', port_string, port_string
            ])
            self.pids['remote_sky_server_port'] = sky_server.port

            port_string = 'tcp:%s' % args.command_port
            subprocess.check_call([
                ADB_PATH, 'forward', port_string, port_string
            ])
            self.pids['remote_sky_command_port'] = args.command_port

        shell_command = self._build_mojo_shell_command(args, is_android)

        # On android we can't launch inside gdb, but rather have to attach.
        if not is_android and args.gdb:
            shell_command = ['gdbserver', ':%d' % GDB_PORT] + shell_command

        print ' '.join(map(pipes.quote, shell_command))
        # This pid is meaningless on android (it's the adb shell pid)
        start_command_pid = subprocess.Popen(shell_command).pid

        if is_android:
            # TODO(eseidel): am start -W does not seem to work?
            pid_tries = 0
            while True:
                pid = self._find_remote_pid_for_package(ANDROID_PACKAGE)
                if pid or pid_tries > 3:
                    break
                logging.debug('No pid for %s yet, waiting' % ANDROID_PACKAGE)
                time.sleep(5)
                pid_tries += 1

            if not pid:
                logging.error('Failed to find mojo_shell pid on device!')
                return
            self.pids['mojo_shell_pid'] = pid
        else:
            self.pids['mojo_shell_pid'] = start_command_pid

        if args.gdb and is_android:
            # We push our own copy of gdbserver with the package since
            # the default gdbserver is a different version from our gdb.
            package_path = \
                self._find_install_location_for_package(ANDROID_PACKAGE)
            gdb_server_path = os.path.join(
                os.path.dirname(package_path), 'lib/arm/gdbserver')
            gdbserver_cmd = [
                ADB_PATH, 'shell',
                gdb_server_path, '--attach',
                ':%d' % GDB_PORT,
                str(self.pids['mojo_shell_pid'])
            ]
            print ' '.join(map(pipes.quote, gdbserver_cmd))
            self.pids['adb_shell_gdbserver_pid'] = \
                subprocess.Popen(gdbserver_cmd).pid

            port_string = 'tcp:%d' % GDB_PORT
            subprocess.check_call([
                ADB_PATH, 'forward', port_string, port_string
            ])
            self.pids['remote_gdbserver_port'] = GDB_PORT

        if not args.gdb:
            if not self._wait_for_sky_command_port():
                logging.error('Failed to start sky')
                self.stop_command(None)
        else:
            # We could just run gdb_attach_command here, but when I do that
            # it auto-suspends in my zsh.  Unclear why.
            # self.gdb_attach_command(args)
            print "Run 'skydb gdb_attach' to attach."

    def _kill_if_exists(self, key, name):
        pid = self.pids.pop(key, None)
        if not pid:
            logging.info('No pid for %s, nothing to do.' % name)
            return
        logging.info('Killing %s (%d).' % (name, pid))
        try:
            os.kill(pid, signal.SIGTERM)
        except OSError:
            logging.info('%s (%d) already gone.' % (name, pid))

    def stop_command(self, args):
        # TODO(eseidel): mojo_shell crashes when attempting graceful shutdown.
        # self._run_basic_command('/quit')

        self._kill_if_exists('sky_server_pid', 'sky_server')

        if 'remote_sky_server_port' in self.pids:
            port_string = 'tcp:%s' % self.pids['remote_sky_server_port']
            subprocess.call([ADB_PATH, 'reverse', '--remove', port_string])

        if 'remote_sky_command_port' in self.pids:
            # adb forward --remove takes the *host* port, not the remote port.
            port_string = 'tcp:%s' % self.pids['sky_command_port']
            subprocess.call([ADB_PATH, 'forward', '--remove', port_string])

            subprocess.call([
                ADB_PATH, 'shell', 'am', 'force-stop', ANDROID_PACKAGE])
        else:
            # Only try to kill mojo_shell if it's running locally.
            self._kill_if_exists('mojo_shell_pid', 'mojo_shell')

        if 'remote_gdbserver_port' in self.pids:
            self._kill_if_exists('adb_shell_gdbserver_pid',
                'adb shell gdbserver')

            port_string = 'tcp:%s' % self.pids['remote_gdbserver_port']
            subprocess.call([ADB_PATH, 'forward', '--remove', port_string])
        self.pids = {} # Clear out our pid file.

    def _url_from_args(self, args):
        if urlparse.urlparse(args.url_or_path).scheme:
            return args.url_or_path
        # The load happens on the remote device, use the remote port.
        remote_sky_server_port = self.pids.get('remote_sky_server_port',
            self.pids['sky_server_port'])
        return SkyServer.url_for_path(remote_sky_server_port,
            self.pids['sky_server_root'], args.url_or_path)

    def load_command(self, args):
        self._run_basic_command('/load', self._url_from_args(args))

    def _read_mojo_map(self):
        # TODO(eseidel): Does not work for android.
        mojo_map_path = "/tmp/mojo_shell.%d.maps" % self.pids['mojo_shell_pid']
        with open(mojo_map_path, 'r') as maps_file:
            lines = maps_file.read().strip().split('\n')
            return dict(map(lambda line: line.split(' '), lines))

    def stop_tracing_command(self, args):
        file_name = args.file_name
        trace = self._send_command_to_sky('/stop_tracing').content
        with open(file_name, "wb") as trace_file:
            trace_file.write(trace)
        print "Trace saved in %s" % file_name

    def stop_profiling_command(self, args):
        self._run_basic_command('/stop_profiling')
        mojo_map = self._read_mojo_map()

        # TODO(eseidel): We should have a helper for resolving urls, etc.
        remote_server_port = self.pids.get('remote_sky_server_port', self.pids['sky_server_port'])
        build_dir_url = SkyServer.url_for_path(
            remote_server_port,
            self.pids['sky_server_root'],
            self.pids['build_dir'])

        # Map /tmp cache paths to urls and then to local build_dir paths.
        def map_to_local_paths(match):
            path = match.group('mojo_path')
            url = mojo_map.get(path)
            if url and url.startswith(build_dir_url):
                return url.replace(build_dir_url, self.pids['build_dir'])
            return match.group(0)

        MOJO_PATH_RE = re.compile(r'(?P<mojo_path>\S+\.mojo)')
        MOJO_NAME_RE = re.compile(r'(?P<mojo_name>\w+)\.mojo')

        with open("sky_viewer.pprof", "rb+") as profile_file:
            # ISO-8859-1 can represent arbitrary binary while still keeping
            # ASCII characters in the ASCII range (allowing us to regexp).
            # http://en.wikipedia.org/wiki/ISO/IEC_8859-1
            as_string = profile_file.read().decode('iso-8859-1')
            # Using the mojo_shell.PID.maps file tmp paths to build_dir paths.
            as_string = MOJO_PATH_RE.sub(map_to_local_paths, as_string)
            # In release foo.mojo is stripped but libfoo_library.so isn't.
            as_string = MOJO_NAME_RE.sub(r'lib\1_library.so', as_string)
            profile_file.seek(0)
            profile_file.write(as_string.encode('iso-8859-1'))
            profile_file.truncate()

    def _command_base_url(self):
        return 'http://localhost:%s' % self.pids['sky_command_port']

    def _send_command_to_sky(self, command_path, payload=None):
        url = 'http://localhost:%s%s' % (
            self.pids['sky_command_port'], command_path)
        if payload:
            response = requests.post(url, payload)
        else:
            response = requests.get(url)
        return response

    def _run_basic_command(self, command_path, payload=None):
        print self._send_command_to_sky(command_path, payload=payload).text

    # FIXME: These could be made into a context object with __enter__/__exit__.
    def _load_pid_file(self, path):
        try:
            with open(path, 'r') as pid_file:
                return json.load(pid_file)
        except:
            if os.path.exists(path):
                logging.warn('Failed to read pid file: %s' % path)
            return {}

    def _write_pid_file(self, path, pids):
        try:
            with open(path, 'w') as pid_file:
                json.dump(pids, pid_file, indent=2, sort_keys=True)
        except:
            logging.warn('Failed to write pid file: %s' % path)

    def _add_basic_command(self, subparsers, name, url_path, help_text):
        parser = subparsers.add_parser(name, help=help_text)
        command = lambda args: self._run_basic_command(url_path)
        parser.set_defaults(func=command)

    def _wait_for_sky_command_port(self):
        tries = 0
        while True:
            try:
                self._run_basic_command('/')
                return True
            except:
                tries += 1
                if tries == 3:
                    logging.warn('Still waiting for sky on port %s' %
                        self.pids['sky_command_port'])
                if tries > 10:
                    return False
                time.sleep(1)

    def logcat_command(self, args):
        TAGS = [
            'AndroidHandler',
            'MojoMain',
            'MojoShellActivity',
            'MojoShellApplication',
            'chromium',
        ]
        subprocess.call([ADB_PATH, 'logcat', '-d', '-s'] + TAGS)

    def _pull_system_libraries(self, system_libs_root):
        # Pull down the system libraries this pid has already mapped in.
        # TODO(eseidel): This does not handle dynamic loads.
        library_cacher_path = os.path.join(
            self.paths.sky_tools_directory, 'android_library_cacher.py')
        subprocess.call([
            library_cacher_path, system_libs_root, self.pids['mojo_shell_pid']
        ])

        # TODO(eseidel): adb_gdb does, this, unclear why solib-absolute-prefix
        # doesn't make this explicit listing not necessary?
        return subprocess.check_output([
            'find', system_libs_root,
            '-mindepth', '1',
            '-maxdepth', '4',
            '-type', 'd',
        ]).strip().split('\n')

    def _add_android_library_links(self, links_path):
        # TODO(eseidel): This might not match mojo_shell on the device?
        # TODO(eseidel): Should we pass libmojo_shell.so as 'file' to gdb?
        shell_link_path = os.path.join(links_path, 'libmojo_shell.so')
        if os.path.lexists(shell_link_path):
            os.unlink(shell_link_path)
        os.symlink(self.paths.mojo_shell_path, shell_link_path)

    def gdb_attach_command(self, args):
        self.paths = self._create_paths_for_build_dir(self.pids['build_dir'])

        if not os.path.exists(CACHE_LINKS_PATH):
            os.makedirs(CACHE_LINKS_PATH)
        cache_linker_path = os.path.join(
             self.paths.sky_tools_directory, 'mojo_cache_linker.py')
        subprocess.check_call([
            cache_linker_path, CACHE_LINKS_PATH, self.paths.build_dir])

        symbol_search_paths = [
            self.pids['build_dir'],
            CACHE_LINKS_PATH,
        ]
        gdb_path = '/usr/bin/gdb'

        eval_commands = [
            'directory %s' % self.paths.src_root,
            'file %s' % self.paths.mojo_shell_path,
            'target remote localhost:%s' % GDB_PORT,
        ]

        # A bunch of extra work is needed for android:
        if 'remote_sky_server_port' in self.pids:
            self._add_android_library_links(CACHE_LINKS_PATH)

            system_lib_dirs = self._pull_system_libraries(SYSTEM_LIBS_ROOT_PATH)
            eval_commands.append(
                'set solib-absolute-prefix %s' % SYSTEM_LIBS_ROOT_PATH)

            symbol_search_paths = system_lib_dirs + symbol_search_paths

            # TODO(eseidel): We need to look up the toolchain somehow?
            gdb_path = os.path.join(SRC_ROOT, 'third_party/android_tools/ndk/'
                'toolchains/arm-linux-androideabi-4.9/prebuilt/linux-x86_64/'
                'bin/arm-linux-androideabi-gdb')

        # Set solib-search-path after letting android modify symbol_search_paths
        eval_commands.append(
            'set solib-search-path %s' % ':'.join(symbol_search_paths))

        exec_command = [gdb_path]
        for command in eval_commands:
            exec_command += ['--eval-command', command]

        print " ".join(exec_command)

        # Write out our pid file before we exec ourselves.
        self._write_pid_file(PID_FILE_PATH, self.pids)

        # Exec gdb directly to avoid python intercepting symbols, etc.
        os.execv(exec_command[0], exec_command)

    def print_crash_command(self, args):
        logcat_cmd = [ADB_PATH, 'logcat', '-d']
        logcat = subprocess.Popen(logcat_cmd, stdout=subprocess.PIPE)

        stack_path = os.path.join(SRC_ROOT,
            'tools', 'android_stack_parser', 'stack')
        stack = subprocess.Popen([stack_path, '-'], stdin=logcat.stdout)
        logcat.wait()
        stack.wait()

    def pids_command(self, args):
        print json.dumps(self.pids, indent=1)

    def main(self):
        logging.basicConfig(level=logging.WARNING)
        logging.getLogger("requests").setLevel(logging.WARNING)

        self.pids = self._load_pid_file(PID_FILE_PATH)

        parser = argparse.ArgumentParser(description='Sky launcher/debugger')
        subparsers = parser.add_subparsers(help='sub-command help')

        start_parser = subparsers.add_parser('start',
            help='launch a new mojo_shell with sky')
        start_parser.add_argument('--gdb', action='store_true')
        start_parser.add_argument('--command-port', type=int,
            default=DEFAULT_SKY_COMMAND_PORT)
        start_parser.add_argument('--use-osmesa', action='store_true',
            default=self._in_chromoting())
        start_parser.add_argument('build_dir', type=str)
        start_parser.add_argument('url_or_path', nargs='?', type=str,
            default=DEFAULT_URL)
        start_parser.add_argument('--show-command', action='store_true',
            help='Display the shell command and exit')
        start_parser.set_defaults(func=self.start_command)

        stop_parser = subparsers.add_parser('stop',
            help=('stop sky (as listed in %s)' % PID_FILE_PATH))
        stop_parser.set_defaults(func=self.stop_command)

        pids_parser = subparsers.add_parser('pids',
            help='dump the current skydb pids file')
        pids_parser.set_defaults(func=self.pids_command)

        logcat_parser = subparsers.add_parser('logcat',
            help=('dump sky-related logs from device'))
        logcat_parser.set_defaults(func=self.logcat_command)

        print_crash_parser = subparsers.add_parser('print_crash',
            help=('dump (and symbolicate) recent crash-stacks'))
        print_crash_parser.set_defaults(func=self.print_crash_command)

        gdb_attach_parser = subparsers.add_parser('gdb_attach',
            help='launch gdb and attach to gdbserver launched from start --gdb')
        gdb_attach_parser.set_defaults(func=self.gdb_attach_command)

        self._add_basic_command(subparsers, 'start_tracing', '/start_tracing',
            'starts tracing the running sky instance')
        self._add_basic_command(subparsers, 'reload', '/reload',
            'reload the current page')
        self._add_basic_command(subparsers, 'start_profiling', '/start_profiling',
            'starts profiling the running sky instance (Linux only)')

        stop_tracing_parser = subparsers.add_parser('stop_tracing',
            help='stops tracing the running sky instance')
        stop_tracing_parser.add_argument('file_name', type=str, default='sky_viewer.trace')
        stop_tracing_parser.set_defaults(func=self.stop_tracing_command)

        stop_profiling_parser = subparsers.add_parser('stop_profiling',
            help='stops profiling the running sky instance (Linux only)')
        stop_profiling_parser.set_defaults(func=self.stop_profiling_command)

        load_parser = subparsers.add_parser('load',
            help='load a new page in the currently running sky')
        load_parser.add_argument('url_or_path', type=str)
        load_parser.set_defaults(func=self.load_command)

        args = parser.parse_args()
        args.func(args)

        self._write_pid_file(PID_FILE_PATH, self.pids)


if __name__ == '__main__':
    SkyDebugger().main()