#!/usr/bin/env python3
import argparse
import os
import subprocess
import sys
from pathlib import Path

from fmt import blue, yellow

ROOT = Path(__file__).parent.parent


def header(msg):
    print(blue(msg))


def label(msg):
    print(yellow(msg))


def cmd(*args):
    print(" ".join(args))
    subprocess.run(args, check=True)


def main():
    parser = argparse.ArgumentParser(
        description="""Run tests for the project

If no flags are specified, run all tests.
"""
    )
    parser.add_argument("--compiler", action="store_true")
    parser.add_argument("--runtime", action="store_true")
    parser.add_argument("--filetest", action="store_true")
    parser.add_argument("--python", action="store_true")
    args = parser.parse_args()

    do_all = not any(vars(args).values())
    if do_all:
        for k in vars(args):
            vars(args)[k] = True

    os.chdir(ROOT)

    # Building and testing the compiler
    if args.compiler:
        header("Compiler...")

        label("Format...")
        cmd("cargo", "fmt", "--", "--check")

        label("Test...")
        cmd("cargo", "test", "--quiet")

        if args.filetest:
            label("Build...")
            cmd("cargo", "build", "--release", "--all")
    elif args.filetest:
        header("Build compiler...")
        cmd("cargo", "build", "--release", "--all")

    # Building and testing the runtime
    os.chdir("rt")
    if args.runtime:
        header("Runtime...")

        label("Format...")
        cmd("cargo", "fmt", "--", "--check")

        label("Test...")
        cmd("cargo", "test", "--quiet")

        if args.filetest:
            label("Build...")
            cmd("cargo", "build", "--release", "--all")
    elif args.filetest:
        header("Build runtime...")
        cmd("cargo", "build", "--release", "--all")
    os.chdir(ROOT)

    # Filetests
    if args.filetest:
        header("Filetests...")

        label("Lit...")
        cmd("tools/lit.py", "ivy-examples/", "test/")

    # Linting the Python scripts
    if args.python:
        header("Python...")

        label("Black...")
        cmd("black", "test/", "--check", "--diff", "--color")

        label("isort...")
        cmd("isort", "test/", "--check", "--diff", "--color")


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(e)
        sys.exit(1)