# -*- coding: UTF-8 -*-
from __future__ import absolute_import
from builtins import hex
import os
import sys
import time
import signal
import plistlib
import argparse
import traceback
import subprocess
import threading

from lib.Utility import FileMananger
from lib.logger import logger as DFLLogger

start_timestamp = time.time()
current_cmd = None


def run_on_host(logger, cmd):
    logger.info(cmd)

    global current_cmd
    current_cmd = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        shell=True)

    (so, se) = current_cmd.communicate()
    returncode = current_cmd.returncode

    so = so.decode('utf-8', 'replace')
    se = se.decode('utf-8', 'replace')

    if returncode:
        logger.verbs(so)
        logger.verbs(se)

    current_cmd = None

    return so, se, returncode

def load_plist(path):
    if sys.version > '3':
        with open(path, 'rb') as f:
            return plistlib.load(f)
    else:
        # Factory station still use Python2 for now.
        return plistlib.readPlist(path)


def read_args():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('-o', '--output', help='Indicate output directory, create it if not exists', required=True)
    argparser.add_argument('-d', '--duration', type=int, help='Indicate the test duration in seconds', required=True)

    argparser.add_argument('-u', '--udid', default='', help='Indicate the device UDID')

    argparser.add_argument('-H', '--host_predication', default='', help='Indicate the host side custormized predication string')
    argparser.add_argument('-D', '--device_predication', default='', help='Indicate the device sise custormized predication string')

    argparser.add_argument('-T', '--timeout', default=180, type=int, help='Operation timeout in seconds')

    args = argparser.parse_args()

    return (args.output, args.duration, args.udid.strip(), args.host_predication, args.device_predication, args.timeout)

def collect_device_CFS_system_log(logger, device_udid, output):
    if not detect_factorysupported_service(logger, device_udid):
        return ''

    config_path = os.path.join(os.path.dirname(__file__), "config/device_syslog_collect.plist")
    so, se, rc = run_on_host(logger, '/usr/local/bin/FactorySupportUtil device-utility --config {} -u {}'.format(config_path, device_udid))

    device_log_path = ''
    for line in (so+se).split('\n'):
        if 'returned value:' in line:
            device_log_path = line.split('returned value:')[-1].strip()
            break
    else:
        return ''

    device_locationID = map_udid_to_locationID(logger, device_udid)
    so, se, rc = run_on_host(logger, '/usr/local/bin/FactorySupportSZFileCopy -l {} -d {} -s {}'.format(device_locationID, output, device_log_path))
    log_path = os.path.join(output, device_log_path.lstrip(os.sep))
    if not os.path.exists(log_path):
        log_path = os.path.join(output, os.path.basename(device_log_path))
        if not os.path.exists(log_path):
            return ''

    return log_path

def collect_device_system_log(logger, output, device_udid, preidcation_config):
    try:
        temp_folder = os.path.join(output, 'temp_device_logarchive')
        FileMananger.create_empty_folder(temp_folder)

        logarchive_path = os.path.join(temp_folder, device_udid) + '.logarchive'
        so, se, rc = run_on_host(logger, 'sudo /usr/bin/log collect --device-udid {} --output "{}"'.format(device_udid, logarchive_path))
        if rc:
            logarchive_path = collect_device_CFS_system_log(logger, device_udid, temp_folder)

        if logarchive_path:
            for log_item in preidcation_config:
                device_prediction = preidcation_config.get(log_item, {}).get('Device', '')
                if device_prediction:
                    so, se, rc = run_on_host(logger, "/usr/local/bin/logdump --predicate '{}' --archive '{}'".format(device_prediction, logarchive_path))
                    if rc:
                        logger.error("Failed to extract {} devicesystem log".format(log_item))
                    else:
                        with open(os.path.join(output, "DFL_device_systemlog_{}.log").format(log_item), 'wb') as f:
                            f.write(so.encode('utf-8'))

    except Exception as e:
        logger.error("{}".format(e))
        logger.verbs(traceback.format_exc())
    finally:
        FileMananger.remove_path(temp_folder)

def collect_host_system_log(logger, output, duration, preidcation_config, device_udid=''):
    try:
        for log_item in preidcation_config:
            host_predication = preidcation_config.get(log_item, {}).get('Host', '')
            test_duration = preidcation_config.get(log_item, {}).get('Duration')
            if not test_duration:
                test_duration = duration

            if host_predication:
                so, se, rc = run_on_host(logger, "sudo /usr/bin/log show --backtrace --debug --info --loss --signpost --last {}s --predicate '{}'".format(get_updated_duration(test_duration), host_predication))
                if rc:
                    logger.error("Failed to show last {}s {} host system log".format(get_updated_duration(test_duration), log_item))
                else:
                    with open(os.path.join(output, "DFL_host_systemlog_{}.log").format(log_item), 'wb') as f:
                        f.write(so.encode('utf-8'))
    except Exception as e:
        logger.error("{}".format(e))
        logger.verbs(traceback.format_exc())


def probe_remotctl_path(logger):
    remotctl_paths = ['/usr/local/bin/remotectl', '/usr/libexec/remotectl']
    for path in remotctl_paths:
        if os.path.exists(path):
            return path

    logger.error("Can't find remotctl at {}".format(remotctl_paths))

    return ''

def map_udid_to_locationID(logger, device_udid):
    ret = None

    so, se, rc = run_on_host(logger, '{} get-property {} LocationID'.format(probe_remotctl_path(logger), device_udid))
    if rc:
        logger.error("Failed to get locationID.")
    else:
        ret = hex(int(so.strip()))

    return ret

def detect_factorysupported_service(logger, device_udid):
    ret = False

    so, se, rc = run_on_host(logger, '{} show {}'.format(probe_remotctl_path(logger), device_udid))
    if rc:
        logger.error("Failed to show device at {}.".format(device_udid))
    else:
        ret = "com.apple.factorysupportd" in so

    return ret

def collect_device_crashlogs(logger, output, device_udid):
    try:
        device_locationID = map_udid_to_locationID(logger, device_udid)
        if not device_locationID:
            logger.error("Failed to get device locationID for udid {}.".format(device_udid))
        elif not detect_factorysupported_service(logger, device_udid):
            logger.error("Doesn't find factorysupported service on device {}".format(device_udid))
        else:
            local_log_folder = os.path.join(output, 'device_crash_logs')
            FileMananger.create_empty_folder(local_log_folder)
            # We don't want to check the target device type, keep the code logic as simple as possible.
            # Try to collect crash logs from embedded(iOS,watchOS...) device
            so, se, rc = run_on_host(logger, '/usr/local/bin/FactorySupportSZFileCopy -l {} -d {} -s /var/mobile/Library/Logs/CrashReporter'.format(device_locationID, local_log_folder))
            # Try to collect crash logs from macOS device
            so, se, rc = run_on_host(logger, '/usr/local/bin/FactorySupportSZFileCopy -l {} -d {} -s /Library/Logs/DiagnosticReports/'.format(device_locationID, local_log_folder))

    except Exception as e:
        logger.error("{}".format(e))
        logger.verbs(traceback.format_exc())


def collect_host_crashlogs(logger, output, duration):
    try:
        local_log_folder = os.path.join(output, 'host_crash_logs')
        FileMananger.create_empty_folder(local_log_folder)

        all_files = FileMananger.list_file_in_folder('/Library/Logs/DiagnosticReports/')
        for file in all_files:
            if os.path.getctime(file) > (time.time() - get_updated_duration(duration)):
                run_on_host(logger, "sudo /bin/cp '{}' '{}'".format(file, local_log_folder))
    except Exception as e:
        logger.error("{}".format(e))
        logger.verbs(traceback.format_exc())

def collect_host_ioreg(logger, output):
    try:
        run_on_host(logger, "sudo /usr/sbin/ioreg -lw0 > {}".format(os.path.join(output, 'host_ioreg.txt')))
    except Exception as e:
        logger.error("{}".format(e))
        logger.verbs(traceback.format_exc())

def get_updated_duration(duration):
    return int(time.time() - start_timestamp + duration)


def main():
    (args_output, args_duration, args_udid, args_host_predication, args_device_predication, args_timeout) = read_args()

    def timer_callback():
        logger.error("{}s timeout...".format(args_timeout))
        if args_udid:
            device_logarchive = os.path.join(args_output, args_udid) + '.logarchive'
            if os.path.exists(device_logarchive):
                FileMananger.remove_path(device_logarchive)
        if current_cmd:
            current_cmd.terminate()
        os.kill(os.getpid(), signal.SIGKILL)

    FileMananger.create_empty_folder(args_output)

    logger = DFLLogger()
    logger.get_logger(os.path.join(args_output, "DFL_Debug.log"))

    timer = threading.Timer(args_timeout, timer_callback)
    timer.start()

    try:
        preidcation_config = load_plist(os.path.join(os.path.dirname(__file__), "config/log_predication.plist"))

        preidcation_config['Station'] = {}
        if args_host_predication:
            preidcation_config['Station']['Host'] = args_host_predication
        if args_device_predication:
            preidcation_config['Station']['Device'] = args_device_predication

        collect_host_system_log(logger, args_output, args_duration, preidcation_config, args_udid)
        collect_host_crashlogs(logger, args_output, args_duration)
        collect_host_ioreg(logger, args_output)
        if args_udid:
            collect_device_system_log(logger, args_output, args_udid, preidcation_config)
            collect_device_crashlogs(logger, args_output, args_udid)
        else:
            logger.info("Skip device log collecting because udid is empty.")
    except:
        logger.error(traceback.format_exc())

    timer.cancel()


if __name__ == "__main__":
    main()