#!/usr/bin/env python3

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

# Set emacs mode to Python
# -*- mode: Python;-*-

"""Rebuild and save cbmc-viewer reports using litani proof artifacts.

The run-cbmc-proofs.py script uses litani to run cbmc and cbmc-viewer
to build proof reports.  This script assumes that run-cbmc-proofs.py
has been run in the current directory, and uses the litani artifacts
produced by the script to rebuild the cbmc-viewer reports.  This
script then saves the cbmc-viewer reports by copying them into a
directory specified on the command line.
"""

from pathlib import Path
import json
import logging
import shutil
import subprocess

import ninja

import arguments

################################################################
# Command line arguments

def parse_arguments():
    """Parse command line arguments"""

    description= 'Rebuild and save cbmc-viewer reports using litani proof artifacts.'
    options = [
        {'flag': 'results',
         'help': 'Copy results into RESULTS.'},
        {'flag': '--copy-only',
         'action': 'store_true',
         'help': 'Copy results without generating results'},
    ]
    args = arguments.create_parser(options, description).parse_args()
    arguments.configure_logging(args)
    return args

################################################################
# The litani cache
#
# When litani runs, it caches everything there is to know about the
# run in a file named 'run.json'. It writes the path to the cache
# directory containing 'run.json' to a file named '.litani_cache_dir'
# in the current working directory.
#
# The file run.json contains the cbmc-viewer commands that this script
# should re-run.  The file is huge, but stripped down to the data we
# care about it looks like
#
# {
#   "pipelines": [
#     {
#       "ci_stages": [
#         {
#           "jobs": [
#             {
#               "wrapper_arguments": {
#                 "ci_stage": "report",
#                 "command": COMMAND,
#                 "outputs": [ REPORT ],
#             }
#           }
#         ]
#       }
#       ]
#     }
#   ]
# }
#
# There is one pipeline for each proof.  There are multiple stages for each
# pipeline.  There are multiple jobs for each stage.  And there are
# wrapper arguments for each job that record the stage name, the command,
# and output produced by the command.
#
# Each report stage consists of a single job that includes a single
# command and a single output.  The command is the invocation of
# cbmc-viewer and the output is the directory containing the report
# produced by cbmc-viewer.

RUN_JSON = 'run.json'
LITANI_CACHE_DIR = '.litani_cache_dir'
COMMAND = 'command'
REPORT = 'report'

def find_cache():
    """Locate the file run.json in the litani cache."""

    cache = Path(RUN_JSON)
    if cache.is_file():
        return cache

    with open(LITANI_CACHE_DIR, encoding='utf-8') as handle:
        cache_dir = Path(handle.readline().strip())
    cache = cache_dir/RUN_JSON
    if cache.is_file():
        return cache

    raise UserWarning(f"Can't find {RUN_JSON}")

def load_cache(cache):
    """Load the file run.json in the litani cache."""

    with open(cache, encoding='utf-8') as handle:
        return json.load(handle)

def parse_cache(cache):
    """Parse the file run.json in the litani cache."""

    jobs  = [job
             for pipeline in cache['pipelines']
             for ci_stage in pipeline['ci_stages']
             for job in ci_stage['jobs']
             if job['wrapper_arguments']['ci_stage'] == 'report']
    return [{
        COMMAND: job['wrapper_arguments']['command'],
        REPORT: job['wrapper_arguments']['outputs'][0]
    } for job in jobs]

################################################################
# Generate reports using cbmc-viewer commands from the litani cache,
# and save the reports (copy them) for later use.

def delete_viewer_reports(jobs):
    """Delete cbmc-viewer reports."""

    for job in jobs:
        report = job[REPORT]
        logging.info("Removing %s", report)
        shutil.rmtree(report,  ignore_errors=True)

def build_viewer_reports(jobs):
    """Build cbmc-viewer reports."""

    for job in jobs:
        report = job[REPORT]
        command = job[COMMAND]
        logging.info("Generating %s", report)
        logging.debug("Running %s", command)
        result = subprocess.run(command.split(),
                                check=True, capture_output=True, text=True)
        for line in result.stdout.splitlines():
            logging.debug(line)
        for line in result.stderr.splitlines():
            logging.debug(line)

def copy_viewer_reports(jobs, reports):
    """Copy cbmc-viewer reports to the reports directory."""

    cwd = Path.cwd()
    shutil.rmtree(reports, ignore_errors=True)
    for job in jobs:
        report = Path(job[REPORT])
        path = report.relative_to(cwd)
        logging.info("Copying %s to %s", path, reports/path)
        shutil.copytree(path, reports/path, dirs_exist_ok=True)

################################################################
# Run cbmc-viewer commands concurrently with ninja

def build_ninja(jobs, debug=False):
    """Build ninja file for running cbmc-viewer commands"""

    rules, builds = [], []
    for job in jobs:
        report_dir = Path(job[REPORT])
        proof_dir = report_dir.parent
        proof_name = proof_dir.name
        cbmc_viewer = job[COMMAND]

        rules.append({
            "name": proof_name,
            "description": proof_name,
            "command": cbmc_viewer + ("" if debug else " >/dev/null") + " 2>&1"
        })
        builds.append({
            "inputs": [],
            "outputs": str(report_dir),
            "rule": proof_name
        })

    with open("build.ninja", "w", encoding='utf-8') as handle:
        writer = ninja.ninja_syntax.Writer(handle)
        for rule in rules:
            writer.rule(**rule)
            writer.newline()
        for build in builds:
            writer.build(**build)
            writer.newline()

def run_ninja(verbose=False):
    """Run ninja to run cbmc-viewer commands"""

    cmd = ['ninja'] + (['--verbose'] if verbose else [])
    logging.debug("Running %s", ' '.join(cmd))
    subprocess.run(cmd, check=True)


################################################################

def main():
    """Rebuild and save cbmc-viewer reports."""

    args = parse_arguments()
    logging.debug(args)

    jobs = parse_cache(load_cache(find_cache()))

    if not args.copy_only:
        delete_viewer_reports(jobs)
        build_ninja(jobs, args.debug)
        run_ninja(args.verbose)
    copy_viewer_reports(jobs, args.results)

if __name__ == '__main__':
    main()