"""
Support script for adding tests
"""
import click
from zipfile import ZipFile
import os
from typing import List

_common_header = """
import unittest
from gradescope_utils.autograder_utils.decorators import number
from neugs_utils.context_managers import Capturing
from neugs_utils import tier, COMMON_ONE, COMMON_TWO, COMMON_THREE, common_msg, strip_prompts
from neugs_utils.common_tests import check_style


student_files = [] # setup incase there are more than one import section
"""


_common_main = """
if __name__ == '__main__':
    unittest.main()
"""

_import_template = """
## Autogenerated import based on test-add
try:
    import {0}
    student_files += ['{1}']
except ImportError:
    try:
        import {2} as {0}
        student_files += ['{3}']
    except ImportError:
        pass

"""

#class TestStarRating(unittest.TestCase):


def build_import(file : str) -> str: 
    """builds a series of try/except statements based on _import_template.

    Args:
        file (str): the file to import

    Returns:
        str: returns an import statement to work locally 
        and on gs auto-grader
    """
    # first get the file parts
    file = os.path.normpath(file).replace("\\", "/")
    dir = os.path.dirname(file)
    file_only = os.path.basename(file)
    mod_base = os.path.splitext(file_only)[0]
    mod_long = '.'.join(dir.split('/')) + "." + mod_base
    return _import_template.format(mod_base, file_only, mod_long, file)




@click.command()
@click.option("--cl", "-c", default="Test")
@click.argument("out")
@click.argument("modules", nargs=-1)
def new_test(out : str,  cl : str, modules : List[str]) -> None:
    """Builds a new test file. Requires the file to write to,
    it is suggest to added a class, but not required, and
    suggested to add modules which match the student files 
    being tested. Currently defaults to
    a set of imports related to TierMastery grading. 

    Args:
        out (str): file to save to
        cl (str): test class
        modules (List[str]): files to import (usually student files)
    """
    imports = ''
    for file in modules:
        imports += build_import(file)
    contents = _common_header
    contents += imports
    contents += "\n" ## add an extra blank
    contents += "# Double check class name, remember unit tests likes it to start with Test\n"
    contents += f"class {cl}(unittest.TestCase):\n    pass\n\n"

    contents += _common_main
    with open(out, 'w') as f:
        f.write(contents)

if __name__ == "__main__":
    new_test()
