import io
import os
import urllib.request
import urllib.request as ur
import zipfile

# Save the original urlopen before monkey-patching it.
original_urlopen = urllib.request.urlopen

# Resolve paths relative to this script so it works from any cwd.
script_dir = os.path.dirname(os.path.abspath(__file__))
os.chdir(script_dir)

# Use the current local project zip.
candidate_zip_paths = [
    os.path.join(script_dir, 'classification.zip'),
]


def build_flattened_zip_bytes(local_zip_path):
    # Rebuild the zip in memory so we can filter junk files and inject a
    # compatibility shim when needed.
    output = io.BytesIO()
    local_zip_name = os.path.basename(local_zip_path)

    # Read the original zip and rebuild it in memory.
    with zipfile.ZipFile(local_zip_path, 'r') as source_zip:
        with zipfile.ZipFile(output, 'w', zipfile.ZIP_DEFLATED) as flat_zip:
            added_names = set()

            # The submission autograder uses the first zip entry as the extracted
            # project root, so make sure the expected directory entry exists first.
            if local_zip_name == 'classification.zip':
                flat_zip.writestr('classification/', b'')
                added_names.add('classification/')

            # Copy each file into a flattened zip.
            for member in source_zip.infolist():
                
                # Ignore common archive metadata or junk files.
                # Ignore macOS archive metadata folders.
                if member.filename.startswith('__MACOSX/'):
                    continue
                # Ignore macOS Finder metadata files.
                if member.filename.endswith('.DS_Store') or '/.DS_Store' in member.filename:
                    continue
                # Ignore Windows Explorer thumbnail cache files.
                if member.filename.endswith('Thumbs.db') or '/Thumbs.db' in member.filename:
                    continue
                # Ignore Windows folder configuration files.
                if member.filename.endswith('desktop.ini') or '/desktop.ini' in member.filename:
                    continue

                # Keep the original tracking/ project root inside the zip so
                # autograder.py lands directly inside the extracted project directory.
                new_name = member.filename

                # Skip empty names and directory entries after the required root
                # folders have already been created above.
                if not new_name or new_name.endswith('/'):
                    continue

                # Write the file into the in-memory zip.
                flat_zip.writestr(new_name, source_zip.read(member.filename))
                added_names.add(new_name)

            # Python 3.13 removed the stdlib cgi module, but the Berkeley
            # autograder still imports cgi.escape. Add a tiny compatibility shim.
            cgi_targets = ['tracking/cgi.py']
            if local_zip_name == 'classification.zip':
                cgi_targets.append('classification/cgi.py')

            for cgi_target in cgi_targets:
                if cgi_target not in added_names:
                    flat_zip.writestr(
                        cgi_target,
                        "from html import escape as _html_escape\n"
                        "\n"
                        "def escape(s, quote=False):\n"
                        "    return _html_escape(s, quote=quote)\n"
                    )

            imp_targets = ['tracking/imp.py']
            if local_zip_name == 'classification.zip':
                imp_targets.append('classification/imp.py')

            for imp_target in imp_targets:
                if imp_target not in added_names:
                    flat_zip.writestr(
                        imp_target,
                        "import sys\n"
                        "import types\n"
                        "\n"
                        "PY_SOURCE = 1\n"
                        "\n"
                        "def new_module(name):\n"
                        "    module = types.ModuleType(name)\n"
                        "    sys.modules[name] = module\n"
                        "    return module\n"
                        "\n"
                        "def load_module(name, file, filename, details=None):\n"
                        "    module = types.ModuleType(name)\n"
                        "    module.__file__ = filename\n"
                        "    sys.modules[name] = module\n"
                        "    source = file.read()\n"
                        "    exec(compile(source, filename, 'exec'), module.__dict__)\n"
                        "    return module\n"
                    )

    # Rewind the in-memory file before returning it.
    output.seek(0)
    return output


def fake_urlopen(url, *args, **kwargs):
    # url may be either a string or a Request object.
    request_url = getattr(url, 'full_url', url)

    # Print the requested URL for debugging.
    print('urlopen URL:', request_url)

    # Intercept the autograder download and serve a local copy instead.
    if (
        'tracking.zip' in request_url
        or 'cs343-6-classification.zip' in request_url
        or (
            'rpl.cs.utexas.edu' in request_url
            and '/projects/classification/' in request_url
            and request_url.endswith('.zip')
        )
    ):
        print('>>> Intercepting! Returning local zip instead...')

        # Use the first local zip file that actually exists.
        for local_zip_path in candidate_zip_paths:
            if os.path.exists(local_zip_path):
                return build_flattened_zip_bytes(local_zip_path)

        # Fail clearly if no local zip file exists.
        raise FileNotFoundError(
            'Could not find a local autograder zip file. '
            'Expected one of: {}'.format(candidate_zip_paths)
        )

    # Fall back to the original urlopen for every other request.
    return original_urlopen(url, *args, **kwargs)


# Patch both imported names used by urllib.request.
urllib.request.urlopen = fake_urlopen
ur.urlopen = fake_urlopen

# Run the local submission autograder after patching urlopen.
exec(open(os.path.join(script_dir, 'submission_autograder.py')).read())
