import argparse
import atexit
import importlib
import os
import subprocess
import sys
from pathlib import Path

from avl import Coverage

from ._coverage import COVERPACKAGES
from ._elf import EXTENSIONS, parse_elf
from ._isa import ISA, extract_isa
from ._spike import parse_spike_log
from ._trace import TRACE

# Path of the current script (__main__.py) and reference json file
script_path = os.path.dirname(os.path.abspath(__file__))
ref_path = os.path.join(script_path, "instr_dict.json")
packages_path = os.path.join(script_path, "coverage_packages")
def main():

    parser = argparse.ArgumentParser()
    parser.add_argument("--ref", type=str, default=ref_path, help=f"Pointer to instr_dict.json (generated from https://github.com/riscv/riscv-opcodes) (default: {ref_path})")
    parser.add_argument("--extensions", nargs="*", default=None, help="List of extensions")
    parser.add_argument("--packages", nargs="*", default=[packages_path], help="List of python packages containing covergroups")
    parser.add_argument("--elf", default=None, help="Elf file for testcase")
    parser.add_argument("--trace", default=None, help="List of trace files")
    parser.add_argument("--output", default="output", help="Output directory for report")
    parser.add_argument("--html", action="store_true", help="Generate HTML report")
    parser.add_argument("--verbose", action="store_true", help="Verbose messages")

    args = parser.parse_args()

    # Parse ISA
    # Creates reference dict for all encodings supported
    if args.verbose:
        print(f"Parsing reference json file : {args.ref} ... ", end="")

    if args.ref is None or not os.path.exists(args.ref):
        print(f"ERROR : must supply a valid instruction reference json file (searched for: {args.ref})")
        sys.exit(1)
    else:
        extract_isa(args.ref)

    if args.verbose:
        print("OK")

    # Parse the elf file
    # Constructs the dicts of instructions and extensions in the elf
    if args.verbose:
        print(f"Parsing elf file : {args.elf} ... ", end="")

    if args.elf is None or not os.path.exists(args.elf):
        print(f"ERROR : must supply a valid elf file (searched for {args.elf})")
        sys.exit(1)
    else:
        parse_elf(args.elf)

    if args.verbose:
        print("OK")

    # Prune un-supported extensions
    if args.verbose:
        print("Pruning un-supported encodings ... ", end="")

    prune = []
    for k, v in ISA.items():
        # Remove instructions with different base
        if v.base not in ["RV", EXTENSIONS["base"]]:
            prune.append(k)
            continue

        # Remove instructions from unsupported extensions
        matches = list(set(v.extensions) & set(EXTENSIONS.keys()))
        if len(matches) == 0:
            prune.append(k)

    a = len(ISA)
    for p in prune:
        del ISA[p]
    b = len(ISA)

    if args.verbose:
        print(f"{a-b} removed")

    # Parse Tracefile
    if args.verbose:
        print(f"Parsing trace file : {args.trace} ... ", end="")

    if args.trace is None or not os.path.exists(args.trace):
        print(f"ERROR : must supply a valid trace (searched for: {args.trace})")
        sys.exit(1)
    else:
        parse_spike_log(args.trace)

    if args.verbose:
        print("OK")


    # Create output directory
    if args.verbose:
        print(f"Creating output directory : {args.output} ... ", end="")

    try:
        path = Path(args.output)
        path.mkdir(parents=True, exist_ok=True)
    except Exception:
        print("Failed to create {args.output} {e}")
        sys.exit(1)

    if args.verbose:
        print("OK")

    # Create covergroups
    if args.packages is not None:
        if args.verbose:
            print(f"Adding coverage packages : {','.join(args.packages)} ... ", end="")

        for p in map(Path, args.packages):
            assert p.is_dir(), f"Packages must be directories: {p}"

            for f in sorted(p.rglob("*.py")):
                package_name = os.path.basename(f).split(".",1)[0]
                package_spec = importlib.util.spec_from_file_location(package_name, Path(f))
                package_module = importlib.util.module_from_spec(package_spec)
                package_spec.loader.exec_module(package_module)

        if args.verbose:
            print("OK")

    # Process traces
    if args.verbose:
        print("Processing Trace ... ", end="")

    instr = TRACE[0]
    while True:
        if instr is None:
            break

        # Sample all Coverpackages
        for p in COVERPACKAGES:
            p.sample(instr)

        instr = instr.next

    if args.verbose:
        print("OK")

    # Move coverage.json to output
    if args.verbose:
        print(f"Moving coverage.json to {args.output} ... ", end="")

    try:
        # Call the atexit to generate the report - then prevent pre-registered callback
        Coverage().__at_exit__()
        atexit.unregister(Coverage().__at_exit__)

        Path("coverage.json").replace(Path(os.path.join(args.output, "coverage.json")))
    except Exception as e:
        print(f"Failed to relocate coverage.json to {args.output} : {e}")
        sys.exit(1)

    if args.verbose:
        print("OK")

    # Generate report
    if args.html:
        if args.verbose:
            print(f"Generating report {args.output}/html/index.html ... ", end="")

        try:
            subprocess.run(
                ['avl-coverage-analysis', '--path', args.output, "--output", os.path.join(args.output, "html"), "--stats"],
                capture_output=True, text=True, check=True
            )
        except (subprocess.CalledProcessError, FileNotFoundError) as e:
            print(f"Could not run avl-coverage-analysis: {e}")
            sys.exit(1)

        if args.verbose:
            print("OK")

if __name__ == "__main__":
    main()
