"""
functions and utilities for handling zipped protocol bundles
"""
from datetime import date
import json
from pathlib import PurePosixPath, PurePath
from typing import Dict, BinaryIO, TYPE_CHECKING
from zipfile import ZipFile

from opentrons.calibration_storage.helpers import uri_from_definition

from .types import BundleContents

if TYPE_CHECKING:
    from opentrons_shared_data.labware.dev_types import LabwareDefinition

MAIN_PROTOCOL_FILENAME = "protocol.ot2.py"
LABWARE_DIR = "labware"
DATA_DIR = "data"


def _has_files_at_root(zipFile):
    for zipInfo in zipFile.infolist():
        if zipInfo.filename.count("/") == 0:
            return True
    return False


def extract_bundle(bundle: ZipFile) -> BundleContents:  # noqa: C901
    """Extract a bundle and verify its contents and structure."""
    if not _has_files_at_root(bundle):
        raise RuntimeError(
            "No files found in ZIP file's root directory. When selecting "
            "files to zip, make sure to directly select the files "
            "themselves. Do not select their parent directory, which would "
            "result in nesting all files inside that directory in the ZIP."
        )
    try:
        with bundle.open(MAIN_PROTOCOL_FILENAME, "r") as protocol_file:
            py_protocol = protocol_file.read().decode("utf-8")
    except KeyError:
        raise RuntimeError(
            f"Bundled protocol should have a {MAIN_PROTOCOL_FILENAME} "
            + "file in the root directory"
        )
    bundled_labware: Dict[str, "LabwareDefinition"] = {}
    bundled_data = {}
    bundled_python = {}
    for zipInfo in bundle.infolist():
        filepath = PurePosixPath(zipInfo.filename)
        rootpath = filepath.parts[0]

        # skip directories and weird OS-added directories
        # (note: the __MACOSX dir would contain '__MACOSX/foo.py'
        # and other files. This would break our inferences, so we need
        # to exclude all contents of that directory)
        if rootpath == "__MACOSX" or zipInfo.is_dir():
            continue

        with bundle.open(zipInfo) as f:
            if rootpath == LABWARE_DIR and filepath.suffix == ".json":
                labware_def = json.loads(f.read().decode("utf-8"))
                labware_key = uri_from_definition(labware_def)
                if labware_key in bundled_labware:
                    raise RuntimeError(f"Conflicting labware in bundle: {labware_key}")
                bundled_labware[labware_key] = labware_def
            elif rootpath == DATA_DIR:
                # note: data files are read as binary
                bundled_data[str(filepath.relative_to(DATA_DIR))] = f.read()
            elif filepath.suffix == ".py" and str(filepath) != MAIN_PROTOCOL_FILENAME:
                bundled_python[str(filepath)] = f.read().decode("utf-8")

    if not bundled_labware:
        raise RuntimeError("No labware definitions found in bundle.")

    return BundleContents(py_protocol, bundled_labware, bundled_data, bundled_python)


def create_bundle(contents: BundleContents, into_file: BinaryIO):
    """Create a bundle from assumed-good contents"""
    with ZipFile(into_file, mode="w") as zf:
        zf.writestr(MAIN_PROTOCOL_FILENAME, contents.protocol)
        for dataname, datafile in contents.bundled_data.items():
            name = PurePath(dataname).name
            zf.writestr(f"{DATA_DIR}/{name}", datafile)
        for lwdef in contents.bundled_labware.values():
            zipsafe = uri_from_definition(lwdef, "-")
            zf.writestr(f"{LABWARE_DIR}/{zipsafe}.json", json.dumps(lwdef))
        zf.writestr(".bundle_beta", str(date.today()))
