#!/usr/bin/env python3
"""Check an Advanced AI lab submission before you upload it.

Usage:

    python3 check_labs.py Labs_Firstname_Lastname_StudentID.zip

It checks three things and changes nothing:

- the zip is named Labs_Firstname_Lastname_StudentID.zip
- it holds the six graded notebooks, each under the name Colab downloads it with
- each notebook is valid and was saved after it was run

A notebook is marked on what it printed, so one saved without outputs has
nothing to mark.

Needs Python 3.9 or later and nothing else.
"""

from __future__ import annotations

import json
import re
import sys
import zipfile
from pathlib import PurePosixPath

# The graded notebooks, one per week from week 1 to week 6, under the names
# Colab downloads them with. tests/test_labs.py keeps this equal to the
# `graded:` entries in labs.yml.
EXPECTED = (
    "2_1_Supervised_Learning",
    "6_2_Gradient_Descent",
    "7_2_Backpropagation",
    "10_1_1D_Convolution",
    "12_1_Self_Attention",
    "19_4_Temporal_Difference_Methods",
)

# Where the email copy goes, and the subject line it carries.
EMAIL = "eoin+advanced-ai-labs@eoin.ai"
SUBJECT = "Advanced AI labs: {first} {last}, {id}"

# A name is letters, joined by a hyphen or an apostrophe where the name has
# one: Mary-Kate, O'Brien, Seán. A surname with a space in it takes a hyphen
# instead: Mac-Giolla.
_NAME = r"[^\W\d_]+(?:['-][^\W\d_]+)*"
ZIP_NAME = re.compile(
    rf"Labs_(?P<first>{_NAME})_(?P<last>{_NAME})_(?P<id>\d{{8}})\.zip"
)

# Files that macOS, Windows and Jupyter add without being asked.
_JUNK_DIRS = {"__MACOSX", ".ipynb_checkpoints"}
_JUNK_FILES = {".DS_Store", "Thumbs.db", "desktop.ini"}


def is_junk(path: PurePosixPath) -> bool:
    return bool(_JUNK_DIRS.intersection(path.parts)) or path.name in _JUNK_FILES


def basename(path: str) -> str:
    return PurePosixPath(path.replace("\\", "/")).name


def check_name(filename: str) -> list[str]:
    if ZIP_NAME.fullmatch(filename):
        return []
    hint = ""
    if not filename.lower().endswith(".zip"):
        hint = " It is not a .zip file."
    elif not filename.startswith("Labs_"):
        hint = " It must start with Labs_, with a capital L."
    elif " " in filename:
        hint = " It contains a space; use underscores between the parts."
    return [
        f"The file is named {filename!r}. It must be "
        "Labs_Firstname_Lastname_StudentID.zip, for example "
        f"Labs_Aoife_Murphy_24123456.zip, with your 8-digit student ID.{hint}"
    ]


def suggest(stem: str) -> str:
    """The expected name a near miss was probably meant to be, if any."""
    cleaned = re.sub(r"^Copy of ", "", stem)
    cleaned = re.sub(r"\s*\(\d+\)$", "", cleaned)
    for name in EXPECTED:
        if cleaned.lower() == name.lower():
            return name
    return ""


def has_outputs(notebook: dict) -> bool:
    return any(
        cell.get("outputs")
        for cell in notebook.get("cells", [])
        if cell.get("cell_type") == "code"
    )


def check_contents(archive: zipfile.ZipFile) -> list[str]:
    errors: list[str] = []
    files = [
        PurePosixPath(info.filename)
        for info in archive.infolist()
        if not info.is_dir() and not is_junk(PurePosixPath(info.filename))
    ]

    # Accept the notebooks at the top of the zip or inside one folder, because
    # students zip the folder about as often as the files.
    parents = {path.parent for path in files}
    if len(parents) > 1:
        errors.append(
            "The notebooks are spread across more than one folder. Put all six "
            "in the same place: the top of the zip, or one folder."
        )

    found = {path.stem: path for path in files if path.suffix == ".ipynb"}
    for path in files:
        if path.suffix != ".ipynb":
            errors.append(
                f"{path} is not a notebook. Submit the six notebooks and nothing else."
            )

    for stem, path in sorted(found.items()):
        if stem in EXPECTED:
            continue
        meant = suggest(stem)
        if meant in found:
            errors.append(f"{path} is a second copy of {meant}.ipynb. Leave it out.")
        elif meant:
            errors.append(f"{path} must be renamed to {meant}.ipynb.")
        else:
            errors.append(f"{path} is not one of the graded notebooks. Leave it out.")

    for name in EXPECTED:
        if name not in found:
            if not any(suggest(stem) == name for stem in found):
                errors.append(f"{name}.ipynb is missing.")
            continue
        try:
            notebook = json.loads(archive.read(str(found[name])).decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError):
            errors.append(f"{name}.ipynb is not a valid notebook file.")
            continue
        if not isinstance(notebook, dict) or "cells" not in notebook:
            errors.append(f"{name}.ipynb is not a valid notebook file.")
        elif not has_outputs(notebook):
            errors.append(
                f"{name}.ipynb has no outputs. Run every cell, then download it again."
            )
    return errors


def check(path: str) -> list[str]:
    errors = check_name(basename(path))
    try:
        with zipfile.ZipFile(path) as archive:
            errors += check_contents(archive)
    except FileNotFoundError:
        errors.append(f"There is no file at {path}.")
    except zipfile.BadZipFile:
        errors.append(f"{path} is not a zip file, or it is damaged.")
    return errors


def email_subject(path: str) -> str:
    """The subject line for this zip, filled in from its name."""
    match = ZIP_NAME.fullmatch(basename(path))
    assert match is not None, "only called once the name has been checked"
    return SUBJECT.format(**match.groupdict())


def main(argv: list[str]) -> int:
    if len(argv) != 2:
        print("Usage: python3 check_labs.py Labs_Firstname_Lastname_StudentID.zip")
        return 2
    errors = check(argv[1])
    if errors:
        print(f"{len(errors)} problem(s) with {argv[1]}:\n")
        for error in errors:
            print(f"  - {error}")
        print("\nFix these and run the check again before you submit.")
        return 1
    print(
        f"{argv[1]} is ready to submit: all {len(EXPECTED)} notebooks are there and have outputs."
    )
    print(f"\nUpload it to Brightspace, then email the same zip to {EMAIL}")
    print(f"with the subject: {email_subject(argv[1])}")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))