#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright 2010 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Script to enter the ChromeOS chroot with mounted sources.

This script enters the chroot with mounted sources.
"""


__author__ = "asharif@google.com (Ahmad Sharif)"

import argparse
import getpass
import os
import pwd
import sys

from cros_utils import command_executer
from cros_utils import logger
from cros_utils import misc


class MountPoint(object):
    """Mount point class"""

    def __init__(self, external_dir, mount_dir, owner, options=None):
        self.external_dir = os.path.realpath(external_dir)
        self.mount_dir = os.path.realpath(mount_dir)
        self.owner = owner
        self.options = options

    def CreateAndOwnDir(self, dir_name):
        retv = 0
        if not os.path.exists(dir_name):
            command = "mkdir -p " + dir_name
            command += " || sudo mkdir -p " + dir_name
            retv = command_executer.GetCommandExecuter().RunCommand(command)
        if retv != 0:
            return retv
        pw = pwd.getpwnam(self.owner)
        if os.stat(dir_name).st_uid != pw.pw_uid:
            command = "sudo chown -f " + self.owner + " " + dir_name
            retv = command_executer.GetCommandExecuter().RunCommand(command)
        return retv

    def DoMount(self):
        ce = command_executer.GetCommandExecuter()
        mount_signature = "%s on %s" % (self.external_dir, self.mount_dir)
        command = "mount"
        retv, out, _ = ce.RunCommandWOutput(command)
        if mount_signature not in out:
            retv = self.CreateAndOwnDir(self.mount_dir)
            logger.GetLogger().LogFatalIf(retv, "Cannot create mount_dir!")
            retv = self.CreateAndOwnDir(self.external_dir)
            logger.GetLogger().LogFatalIf(retv, "Cannot create external_dir!")
            retv = self.MountDir()
            logger.GetLogger().LogFatalIf(retv, "Cannot mount!")
            return retv
        else:
            return 0

    def UnMount(self):
        ce = command_executer.GetCommandExecuter()
        return ce.RunCommand("sudo umount %s" % self.mount_dir)

    def MountDir(self):
        command = (
            "sudo mount --bind " + self.external_dir + " " + self.mount_dir
        )
        if self.options == "ro":
            command += " && sudo mount --bind -oremount,ro " + self.mount_dir
        retv = command_executer.GetCommandExecuter().RunCommand(command)
        return retv

    def __str__(self):
        ret = ""
        ret += self.external_dir + "\n"
        ret += self.mount_dir + "\n"
        if self.owner:
            ret += self.owner + "\n"
        if self.options:
            ret += self.options + "\n"
        return ret


def Main(argv, return_output=False):
    """The main function."""

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-c",
        "--chromeos_root",
        dest="chromeos_root",
        default="../..",
        help="ChromeOS root checkout directory.",
    )
    parser.add_argument(
        "-t",
        "--toolchain_root",
        dest="toolchain_root",
        help="Toolchain root directory.",
    )
    parser.add_argument(
        "-o", "--output", dest="output", help="Toolchain output directory"
    )
    parser.add_argument(
        "--sudo",
        dest="sudo",
        action="store_true",
        default=False,
        help="Run the command with sudo.",
    )
    parser.add_argument(
        "-r",
        "--third_party",
        dest="third_party",
        help="The third_party directory to mount.",
    )
    parser.add_argument(
        "-m",
        "--other_mounts",
        dest="other_mounts",
        help="Other mount points in the form: " "dir:mounted_dir:options",
    )
    parser.add_argument(
        "-s",
        "--mount-scripts-only",
        dest="mount_scripts_only",
        action="store_true",
        default=False,
        help="Mount only the scripts dir, and not the sources.",
    )
    parser.add_argument(
        "passthrough_argv",
        nargs="*",
        help="Command to be executed inside the chroot.",
    )

    options = parser.parse_args(argv)

    chromeos_root = options.chromeos_root

    chromeos_root = os.path.expanduser(chromeos_root)
    if options.toolchain_root:
        options.toolchain_root = os.path.expanduser(options.toolchain_root)

    chromeos_root = os.path.abspath(chromeos_root)

    tc_dirs = []
    if options.toolchain_root is None or options.mount_scripts_only:
        m = "toolchain_root not specified. Will not mount toolchain dirs."
        logger.GetLogger().LogWarning(m)
    else:
        tc_dirs = [
            options.toolchain_root + "/google_vendor_src_branch/gcc",
            options.toolchain_root + "/google_vendor_src_branch/binutils",
        ]

    for tc_dir in tc_dirs:
        if not os.path.exists(tc_dir):
            logger.GetLogger().LogError(
                "toolchain path " + tc_dir + " does not exist!"
            )
            parser.print_help()
            sys.exit(1)

    if not os.path.exists(chromeos_root):
        logger.GetLogger().LogError(
            "chromeos_root " + options.chromeos_root + " does not exist!"
        )
        parser.print_help()
        sys.exit(1)

    if not os.path.exists(chromeos_root + "/src/scripts/build_packages"):
        logger.GetLogger().LogError(
            options.chromeos_root + "/src/scripts/build_packages" " not found!"
        )
        parser.print_help()
        sys.exit(1)

    version_dir = os.path.realpath(
        os.path.expanduser(os.path.dirname(__file__))
    )

    mounted_tc_root = "/usr/local/toolchain_root"
    full_mounted_tc_root = chromeos_root + "/chroot/" + mounted_tc_root
    full_mounted_tc_root = os.path.abspath(full_mounted_tc_root)

    mount_points = []
    for tc_dir in tc_dirs:
        last_dir = misc.GetRoot(tc_dir)[1]
        mount_point = MountPoint(
            tc_dir,
            full_mounted_tc_root + "/" + last_dir,
            getpass.getuser(),
            "ro",
        )
        mount_points.append(mount_point)

    # Add the third_party mount point if it exists
    if options.third_party:
        third_party_dir = options.third_party
        logger.GetLogger().LogFatalIf(
            not os.path.isdir(third_party_dir),
            "--third_party option is not a valid dir.",
        )
    else:
        third_party_dir = os.path.abspath(
            "%s/../../../third_party" % os.path.dirname(__file__)
        )

    if os.path.isdir(third_party_dir):
        mount_point = MountPoint(
            third_party_dir,
            (
                "%s/%s"
                % (full_mounted_tc_root, os.path.basename(third_party_dir))
            ),
            getpass.getuser(),
        )
        mount_points.append(mount_point)

    output = options.output
    if output is None and options.toolchain_root:
        # Mount the output directory at /usr/local/toolchain_root/output
        output = options.toolchain_root + "/output"

    if output:
        mount_points.append(
            MountPoint(
                output, full_mounted_tc_root + "/output", getpass.getuser()
            )
        )

    # Mount the other mount points
    mount_points += CreateMountPointsFromString(
        options.other_mounts, chromeos_root + "/chroot/"
    )

    last_dir = misc.GetRoot(version_dir)[1]

    # Mount the version dir (v14) at /usr/local/toolchain_root/v14
    mount_point = MountPoint(
        version_dir, full_mounted_tc_root + "/" + last_dir, getpass.getuser()
    )
    mount_points.append(mount_point)

    for mount_point in mount_points:
        retv = mount_point.DoMount()
        if retv != 0:
            return retv

    # Finally, create the symlink to build-gcc.
    command = "sudo chown " + getpass.getuser() + " " + full_mounted_tc_root
    retv = command_executer.GetCommandExecuter().RunCommand(command)

    try:
        CreateSymlink(
            last_dir + "/build-gcc", full_mounted_tc_root + "/build-gcc"
        )
        CreateSymlink(
            last_dir + "/build-binutils",
            full_mounted_tc_root + "/build-binutils",
        )
    except Exception as e:
        logger.GetLogger().LogError(str(e))

    # Now call cros_sdk --enter with the rest of the arguments.
    command = "cd %s/src/scripts && cros_sdk --enter" % chromeos_root

    if len(options.passthrough_argv) > 1:
        inner_command = " ".join(options.passthrough_argv[1:])
        inner_command = inner_command.strip()
        if inner_command.startswith("-- "):
            inner_command = inner_command[3:]
        command_file = "tc_enter_chroot.cmd"
        command_file_path = chromeos_root + "/src/scripts/" + command_file
        retv = command_executer.GetCommandExecuter().RunCommand(
            "sudo rm -f " + command_file_path
        )
        if retv != 0:
            return retv
        with open(command_file_path, "w", encoding="utf-8") as f:
            f.write(inner_command)
        logger.GetLogger().LogCmd(inner_command)
        retv = command_executer.GetCommandExecuter().RunCommand(
            "chmod +x " + command_file_path
        )
        if retv != 0:
            return retv

        if options.sudo:
            command += " sudo ./" + command_file
        else:
            command += " ./" + command_file
        retv = command_executer.GetCommandExecuter().RunCommandGeneric(
            command, return_output
        )
        return retv
    else:
        os.chdir("%s/src/scripts" % chromeos_root)
        ce = command_executer.GetCommandExecuter()
        _, out, _ = ce.RunCommandWOutput("which cros_sdk")
        cros_sdk_binary = out.split()[0]
        return os.execv(cros_sdk_binary, ["", "--enter"])


def CreateMountPointsFromString(mount_strings, chroot_dir):
    # String has options in the form dir:mount:options
    mount_points = []
    if not mount_strings:
        return mount_points
    mount_list = mount_strings.split()
    for mount_string in mount_list:
        mount_values = mount_string.split(":")
        external_dir = mount_values[0]
        mount_dir = mount_values[1]
        if len(mount_values) > 2:
            options = mount_values[2]
        else:
            options = None
        mount_point = MountPoint(
            external_dir,
            chroot_dir + "/" + mount_dir,
            getpass.getuser(),
            options,
        )
        mount_points.append(mount_point)
    return mount_points


def CreateSymlink(target, link_name):
    logger.GetLogger().LogFatalIf(
        target.startswith("/"), "Can't create symlink to absolute path!"
    )
    real_from_file = misc.GetRoot(link_name)[0] + "/" + target
    if os.path.realpath(real_from_file) != os.path.realpath(link_name):
        if os.path.exists(link_name):
            command = "rm -rf " + link_name
            command_executer.GetCommandExecuter().RunCommand(command)
        os.symlink(target, link_name)


if __name__ == "__main__":
    retval = Main(sys.argv)
    sys.exit(retval)