import argparse
import importlib.util
import json
import os
import sys



def import_module_from_path(py_path: str, module_name: str = "student_main"):
    if not os.path.exists(py_path):
        raise RuntimeError(f"Python file not found: {py_path}")

    # Make relative imports work (imports inside submission folder)
    submission_dir = os.path.dirname(os.path.abspath(py_path))
    if submission_dir not in sys.path:
        sys.path.insert(0, submission_dir)

    spec = importlib.util.spec_from_file_location(module_name, py_path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Could not load module spec from: {py_path}")

    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)

    return mod


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--submission_dir", required=True)  # kept for compatibility / sanity checks
    ap.add_argument("--script", required=True)          # path to student's main.py
    ap.add_argument("--function", required=True)        # e.g. "evaluate"
    ap.add_argument("--input", required=True)           # input file path
    ap.add_argument("--output", required=True)          # output json path
    args = ap.parse_args()
    
    # Optional: ensure script is inside submission_dir
    subdir = os.path.abspath(args.submission_dir)
    script = os.path.abspath(args.script)
    if not script.startswith(subdir + os.sep):
        raise RuntimeError(f"--script must be inside --submission_dir.\nscript={script}\nsubmission_dir={subdir}")

    mod = import_module_from_path(script, module_name="student_main")

    if not hasattr(mod, args.function):
        raise RuntimeError(f"{args.function}(...) not found in {script}")

    fn = getattr(mod, args.function)

    # Call the student's function. (Still passes a filepath.)
    pred = fn(args.input)

    with open(args.output, "w") as f:
        json.dump(pred, f)


if __name__ == "__main__":
    main()