# Utility for finding a set of logs to use as input for dataset_generator_main
import argparse
import json
import os
import sys
import urllib.request as urllib
from subprocess import check_call
from infra.data_catalog.client import data_rest_api

from base.proto import WriteProtoAsText, ReadProtoAsText
from mined_metric.builder.proto.data_loader_options_pb2 import ChumUriSet
from vision.data.proto.log_finder_options_pb2 import LogFinderOptions, LogFinderOutput

LOG_FINDER_OPTIONS_PATH = 'vehicle/perception/learning/safetynet/proto/default_log_finder_options.pbtxt'
OUTPUT_FILTERED_NAME = "filtered.pbtxt"
OUTPUT_NAME = "entry.pbtxt"

KILOMETERS_TO_MILES = 0.6213712
MINIMUM_KM_PER_RUN = 2.0


def parse_args(argv):
    parser = argparse.ArgumentParser(__doc__)
    parser.add_argument(
        '--output_dir',
        type=str,
        help=('Output directory'),
        default='/tmp/out',
    )
    parser.add_argument(
        '--log_finder_options',
        type=str,
        help=('Path to log finder options pbtxt'),
        default=LOG_FINDER_OPTIONS_PATH,
    )
    parser.add_argument(
        '--filter_critical',
        action="store_true",
        help=(f"If specified, save logs without critical disengements in "
              "output_dir as {OUTPUT_FILTERED_NAME}. Useful for mining false "
              "positive runs"),
    )
    return parser.parse_args(argv)


def main(argv=sys.argv[1:]):
    args = parse_args(argv)
    options = LogFinderOptions()
    ReadProtoAsText(args.log_finder_options, options)
    # Update output paths in log finder options proto
    options.output_dataset_list_pbtxt_fn = os.path.join(args.output_dir, "dataset.pbtxt")
    options.output_html_fn = os.path.join(args.output_dir, "details.pbtxt")
    options.output_pbtxt_fn = os.path.join(args.output_dir, OUTPUT_NAME)
    options.output_run_ids_in_txt_fn = os.path.join(args.output_dir, "run_id.pbtxt")
    options.output_stats_html_fn = os.path.join(args.output_dir, "stats.pbtxt")
    output_options_path = os.path.join(
        args.output_dir, "log_finder_options.pbtxt")
    WriteProtoAsText(output_options_path, options)
    # Run vision log finder
    check_call([
        'vision/data/tools/find_logs',
        '--options_pbtxt',
        output_options_path
    ])
    print("Output saved to {}".format(options.output_pbtxt_fn))

    if args.filter_critical:
        total_autonomous_km = 0
        output_filtered_fn = os.path.join(
            args.output_dir, OUTPUT_FILTERED_NAME)
        output = LogFinderOutput()
        output_filtered = ChumUriSet()
        ReadProtoAsText(options.output_pbtxt_fn, output)
        for entry in output.entry:
            # Try block in case url lookup throws error.
            try:
                disengements = get_disengagements_from_run_id(entry.run_id)
                any_critical = False
                # Check all disengements and only add entry if no critical
                # disengements are found
                for disengement in disengements:
                    if ('reason' in disengement and
                        disengement['reason'] == "CRITICAL" or
                    'final_reason' in disengement and
                        disengement['final_reason'] == "CRITICAL"):
                        any_critical = True
                if not any_critical:
                    autonomous_km = check_data_catalog(data_rest_api.get_odometry(entry.run_id))['autonomous_delta_km']
                    if autonomous_km < MINIMUM_KM_PER_RUN:
                        continue
                    total_autonomous_km += autonomous_km
                    output_filtered.chum_uri.extend([entry.run_id])
                else:
                    print(f"Filtered out log with critical disengagement {entry.run_id}")
            except Exception as e:
                print(e)
        print( output_filtered_fn)
        WriteProtoAsText(output_filtered_fn, output_filtered)
        with open(output_filtered_fn, 'r') as original: data = original.read()
        with open(output_filtered_fn, 'w') as modified: modified.write(
                f"""# Generated with {args.log_finder_options}
# Total autonomous km: {total_autonomous_km}
# Total autonomous miles: {total_autonomous_km * KILOMETERS_TO_MILES}
""" + data)

def check_data_catalog(result):
    if not result['success'] == True:
        raise Exception(result)

    return result['odometry']

def get_disengagements_from_run_id(run_id):
    """
    Lookup all disengagements for a given run id.

    Args:
        run_id: str
            Run id to lookup

    Returns:
        List of json encoded disengagements
    """
    ARGUS_BACKEND_URL = (
        'https://vehicle-events.zooxlabs.com/v0/disengagement?runIdentifierList=' +
        run_id)
    json_file = urllib.urlopen(ARGUS_BACKEND_URL).read()
    json_decode = json.loads(json_file)
    if 'items' not in json_decode: return None
    disengements = json_decode['items']
    return disengements


if __name__ == "__main__":
    main()
