import datetime
import logging
import os
from pathlib import Path
from pprint import pformat
import tarfile
from typing import Dict, List, Sequence
import typing

from olot.constants import (
    ANNOTATION_LAYER_CONTENT_DIGEST,
    ANNOTATION_LAYER_CONTENT_TYPE,
    ANNOTATION_LAYER_CONTENT_INLAYERPATH,
    ANNOTATION_LAYER_CONTENT_NAME,
)
from olot.dockerdist.convert import check_if_oci_layout_contains_docker_manifests, convert_docker_manifests_to_oci
from olot.enums import RemoveOriginals
from olot.modelpack.model_config import Model, ModelConfig, ModelDescriptor, ModelFS, Type
from olot.oci.oci_config import HistoryItem, OCIManifestConfig

from olot.oci.oci_image_index import Manifest, OCIImageIndex, Platform, read_ocilayout_root_index
from olot.oci.oci_image_manifest import OCIImageManifest, ContentDescriptor, create_oci_image_manifest
from olot.oci.oci_image_layout import verify_ocilayout
from olot.oci.oci_common import MediaTypes

from olot.utils.files import LayerStats, handle_remove, tarball_from_file, targz_from_file
from olot.utils.types import compute_hash_of_str

from olot.modelpack import const as modelpack_consts

logger = logging.getLogger(__name__)


def oci_layers_on_top(
        ocilayout: typing.Union[str, os.PathLike],
        model_files: Sequence[os.PathLike],
        modelcard: typing.Union[os.PathLike, None] = None,
        *,
        labels: typing.Union[dict[str, str], None] = None,
        annotations: typing.Union[dict[str, str], None] = None,
        remove_originals: typing.Union[RemoveOriginals, None] = None,
        add_modelpack: typing.Union[bool, None] = None):
    """
    Add contents to an oci-layout directory as new blob layers

    Args:
        ocilayout: The oci-layout directory of the base image.
        model_files: PathLike array to be added as new blob layers.
        modelcard: PathLike of the README.md of the ModelCarD, will be added as the last layer with compression and annotations. If indicated, it shouldn't be part of model_files.
        labels: labels to be added to the OCI Image Config.
        annotations: annotations to be added to the OCI Image Manifest.
        remove_originals: whether to remove the original content files after having added the layers, default: None.
        add_modelpack: whether to add a ModelPack manifest to the multi-arch oci-layout only if not already present, default: None.
    """
    logger.debug("oci_layers_on_top parameters: %s", pformat(locals()))

    if not isinstance(ocilayout, Path):
        ocilayout = Path(ocilayout)
    if remove_originals:
        logger.info("Invoked with remove_originals=%s to delete original contents after adding as a blob layer.", remove_originals.value)

    # Normalize each path in model_files and modelcard to ensure consistency
    model_files = [Path(mf).resolve() for mf in model_files]
    if modelcard is not None:
        modelcard = Path(modelcard).resolve()
    for model in model_files:
        model = Path(model)
        if model.is_dir():
            logger.warning(f"One of the input is a whole directory and will result in non-efficient layer-ing: {model}")

    verify_ocilayout(ocilayout)
    if check_if_oci_layout_contains_docker_manifests(ocilayout):
        logger.warning("OCI layout contains Docker distribution manifests, converting them to OCI format")
        convert_docker_manifests_to_oci(ocilayout)
    ocilayout_root_index: OCIImageIndex = read_ocilayout_root_index(ocilayout)
    ocilayout_indexes: Dict[str, OCIImageIndex] = crawl_ocilayout_indexes(ocilayout, ocilayout_root_index)
    ocilayout_manifests: Dict[str, OCIImageManifest] = crawl_ocilayout_manifests(ocilayout, ocilayout_indexes, ocilayout_root_index)
    new_layers: Dict[str, LayerStats] = {} # layer digest : diff_id

    # check configuration is consistent
    add_modelpack = check_and_sanitize_flag_add_modelpack(add_modelpack, ocilayout_indexes, ocilayout_manifests)
    if modelcard in model_files:
        if remove_originals:
            raise ValueError("ModelCard detected in model_files, while remove_originals flag is set; this is not allowed as it would remove the original ModelCard before having a chance of adding it as its proper layer.")
        else:
            logger.warning("ModelCard detected in model_files, this will result in duplicated layers for the ModelCard (negligible, but not optimal).")

    sha256_path = ocilayout / "blobs" / "sha256"
    for model in model_files:
        model = Path(model)
        new_layer = tarball_from_file(model, sha256_path)
        new_layers[new_layer.layer_digest] = new_layer
        if remove_originals:
            handle_remove(model)
    if modelcard is not None:
        new_layer = targz_from_file(Path(modelcard), sha256_path)
        new_layers[new_layer.layer_digest] = new_layer
        if remove_originals == RemoveOriginals.ALL:
            handle_remove(modelcard)

    new_ocilayout_manifests: Dict[str, str] = {}
    for manifest_hash, manifest in ocilayout_manifests.items():
        logger.debug("manifest_hash: %s, manifest.mediaType: %s", manifest_hash, manifest.mediaType)
        config_sha = manifest.config.digest.removeprefix("sha256:")
        mc = None
        with open(ocilayout / "blobs" / "sha256" / config_sha, "r") as cf:
            mc = OCIManifestConfig.model_validate_json(cf.read())
            if mc.history is None:
                mc.history = []
        for new_layer in new_layers.values():
            layer_digest = new_layer.layer_digest
            layer_stat = os.stat(ocilayout / "blobs" / "sha256" / layer_digest)
            size = layer_stat.st_size
            ctime = layer_stat.st_ctime
            mt = MediaTypes.layer if layer_digest == new_layer.diff_id else MediaTypes.layer_gzip
            la = {"org.opencontainers.image.title": new_layer.title} # if ever used for OCI artifact, add `unpack` annotations 
            is_modelcard = layer_digest != new_layer.diff_id
            if is_modelcard:
                la["io.opendatahub.modelcar.layer.type"] = "modelcard"
            if new_layer.input_hash:
                la[ANNOTATION_LAYER_CONTENT_DIGEST] = "sha256:"+new_layer.input_hash
            la[ANNOTATION_LAYER_CONTENT_TYPE] = new_layer.input_type
            la[ANNOTATION_LAYER_CONTENT_INLAYERPATH] = new_layer.in_layer_path
            la[ANNOTATION_LAYER_CONTENT_NAME] = new_layer.title
            cd = ContentDescriptor(
                mediaType=mt,
                digest="sha256:"+layer_digest,
                size=size,
                urls=None,
                data=None,
                artifactType=None,
                annotations=la
            )
            mc.rootfs.diff_ids.append("sha256:"+new_layer.diff_id)
            hi = HistoryItem(
                created=datetime.datetime.fromtimestamp(ctime, tz=datetime.timezone.utc).isoformat(timespec="seconds").replace("+00:00","Z"),
                created_by="olot oci_layers_on_top "+new_layer.title+(" (modelcard)" if is_modelcard else "")
            )
            mc.history.append(hi)
            manifest.layers.append(cd)
        # add labels to the OCI Image Config
        if labels is not None:
            if mc.config is None:  # should rarely happen, but just in case
                logger.warning("OCI Image Config is missing, will not add labels")
            else:
                if mc.config.Labels is None:
                    mc.config.Labels = {}
                mc.config.Labels.update(labels)
        # save the OCI Image Config
        mc_json = mc.model_dump_json(exclude_none=True)
        with open(ocilayout / "blobs" / "sha256" / config_sha, "w") as cf:
            cf.write(mc_json)
        mc_json_hash = compute_hash_of_str(mc_json)
        os.rename(ocilayout / "blobs" / "sha256" / config_sha, ocilayout / "blobs" / "sha256" / mc_json_hash)
        logger.debug("Renamed config from: %s to %s", config_sha, mc_json_hash)
        config_sha = mc_json_hash
        manifest.config.digest = "sha256:" + config_sha
        manifest.config.size = os.stat(ocilayout / "blobs" / "sha256" / config_sha).st_size
        if manifest.annotations is None:
            manifest.annotations = {}
        manifest.annotations["io.opendatahub.author"] = "olot"
        if modelcard is not None:
            manifest.annotations["io.opendatahub.layers.modelcard"] = "sha256:"+next(reversed(new_layers.keys())) # identify ModelCarD layer from Image Manifest
        if annotations is not None:
            manifest.annotations.update(annotations)
        check_manifest(manifest, mc)
        # save the OCI Image Manifest
        manifest_json = manifest.model_dump_json(exclude_none=True)
        with open(ocilayout / "blobs" / "sha256" / manifest_hash, "w") as cf:
            cf.write(manifest_json)
        manifest_json_hash = compute_hash_of_str(manifest_json)
        os.rename(ocilayout / "blobs" / "sha256" / manifest_hash, ocilayout / "blobs" / "sha256" / manifest_json_hash)
        logger.debug("Renamed Manifest from: %s to %s", manifest_hash, manifest_json_hash)
        new_ocilayout_manifests[manifest_hash] = manifest_json_hash
        manifest_hash = manifest_json_hash
    logger.debug("new_ocilayout_manifests: %s", pformat(new_ocilayout_manifests))
    if add_modelpack:
        modelpack_manifest_hash = add_modelpack_manifest(ocilayout, new_layers)
    
    new_ocilayout_indexes: Dict[str, str] = {}
    for index_hash, index in ocilayout_indexes.items():
        logger.debug("index_hash: %s, index.mediaType: %s", index_hash, index.mediaType)
        for m in index.manifests:
            digest = m.digest.removeprefix("sha256:")
            if digest in new_ocilayout_manifests.keys():
                lookup_new_hash = new_ocilayout_manifests[m.digest.removeprefix("sha256:")]
                logger.info("old manifest %s is now at %s", m.digest, lookup_new_hash)
                m.digest = "sha256:" + lookup_new_hash
                m.size = os.stat(ocilayout / "blobs" / "sha256" / lookup_new_hash).st_size
            else:
                logger.info("manifest %s was unchanged", digest)
        if add_modelpack:
            if modelpack_manifest_hash in [m.digest.removeprefix("sha256:") for m in index.manifests]:
                raise ValueError(f"ModelPack manifest already present in Index {index_hash} while add_modelpack=True")
            else:
                logger.debug("adding ModelPack manifest to the index %s", index_hash)
                index.manifests.append(Manifest(
                    mediaType=MediaTypes.manifest,
                    digest="sha256:"+modelpack_manifest_hash,
                    size=os.stat(ocilayout / "blobs" / "sha256" / modelpack_manifest_hash).st_size,
                    urls=None,
                    platform=Platform( # type: ignore[call-arg]
                        architecture="unknown",
                        os="unknown",
                    ),
                    annotations={"io.opendatahub.author": "olot", "io.opendatahub.modelcar.manifest.type":"modelpack"},
                ))
        index_json = index.model_dump_json(exclude_none=True)
        with open(ocilayout / "blobs" / "sha256" / index_hash, "w") as idxf:
            idxf.write(index_json)
        index_json_hash = compute_hash_of_str(index_json)
        os.rename(ocilayout / "blobs" / "sha256" / index_hash, ocilayout / "blobs" / "sha256" / index_json_hash)
        logger.debug("Renamed Index from: %s to %s", index_hash, index_json_hash)
        new_ocilayout_indexes[index_hash] = index_json_hash
    logger.debug("new_ocilayout_indexes: %s", pformat(new_ocilayout_indexes))
    
    for entry in ocilayout_root_index.manifests:
        if entry.mediaType == MediaTypes.index:
            lookup_new_hash = new_ocilayout_indexes[entry.digest.removeprefix("sha256:")]
            logger.debug("old index %s is now at %s", entry.digest, lookup_new_hash)
            entry.digest = "sha256:" + lookup_new_hash
            entry.size = os.stat(ocilayout / "blobs" / "sha256" / lookup_new_hash).st_size
        elif entry.mediaType == MediaTypes.manifest:
            digest = entry.digest.removeprefix("sha256:")
            if digest in new_ocilayout_manifests.keys():
                lookup_new_hash = new_ocilayout_manifests[entry.digest.removeprefix("sha256:")]
                logger.info("old manifest %s is now at %s", entry.digest, lookup_new_hash)
                entry.digest = "sha256:" + lookup_new_hash
                entry.size = os.stat(ocilayout / "blobs" / "sha256" / lookup_new_hash).st_size
            else:
                logger.info("manifest %s was unchanged", digest)
        else:
            raise ValueError(f"unknown root index mediaType {entry.mediaType}")
    if add_modelpack:
        if modelpack_manifest_hash in [m.digest.removeprefix("sha256:") for m in ocilayout_root_index.manifests]:
            raise ValueError("ModelPack manifest already present in Root index in oci-layout, while add_modelpack=True")
        else:
            logger.debug("adding ModelPack manifest to the root index")
            ocilayout_root_index.manifests.append(Manifest(
                mediaType=MediaTypes.manifest,
                digest="sha256:"+modelpack_manifest_hash,
                size=os.stat(ocilayout / "blobs" / "sha256" / modelpack_manifest_hash).st_size,
                urls=None,
                platform=Platform( # type: ignore[call-arg]
                    architecture="unknown",
                    os="unknown",
                ),
                annotations={"io.opendatahub.author": "olot", "io.opendatahub.modelcar.manifest.type":"modelpack"},
            ))
    with open(ocilayout / "index.json", "w") as root_idx_f:
        root_idx_f.write(ocilayout_root_index.model_dump_json(exclude_none=True))


def add_modelpack_manifest(ocilayout: Path, new_layers: Dict[str, LayerStats]) -> str:
    """add a ModelPack manifest to the oci-layout
    """
    model_config = Model(
        descriptor=ModelDescriptor(name=None), # eventually config and metadata will be provided programmatically
        modelfs=ModelFS(type=Type.layers, diffIds=list(b.diff_id for b in new_layers.values())),
        config=ModelConfig(),
    )
    model_config_json = model_config.model_dump_json(exclude_none=True)
    model_config_hash = compute_hash_of_str(model_config_json)
    with open(ocilayout / "blobs" / "sha256" / model_config_hash, "w") as mpcf:
        mpcf.write(model_config_json)
    cds = []
    for layer in new_layers.values():
        layer_digest = layer.layer_digest
        layer_stat = os.stat(ocilayout / "blobs" / "sha256" / layer_digest)
        size = layer_stat.st_size
        la = {"org.opencontainers.image.title": layer.title}
        is_modelcard = layer_digest != layer.diff_id
        mt = modelpack_consts.MEDIATYPEMODELDOCGZIP if is_modelcard else modelpack_consts.MEDIATYPEMODELWEIGHT
        cd = ContentDescriptor(
            mediaType=mt,
            digest="sha256:"+layer_digest,
            size=size,
            urls=None,
            data=None,
            artifactType=None,
            annotations=la
        )
        cds.append(cd)
    manifest = create_oci_image_manifest(
        mediaType=MediaTypes.manifest, # default already, but for expliciting here Modelpack spec adherence
        artifactType=modelpack_consts.ARTIFACTTYPEMODELMANIFEST,
        config=ContentDescriptor(
            mediaType=modelpack_consts.MEDIATYPEMODELCONFIG,
            artifactType=None,
            digest="sha256:"+model_config_hash,
            size=os.stat(ocilayout / "blobs" / "sha256" / model_config_hash).st_size,
            urls=None,
            data=None,
        ),
        layers=cds
    )
    manifest_json = manifest.model_dump_json(exclude_none=True)
    manifest_hash = compute_hash_of_str(manifest_json)
    with open(ocilayout / "blobs" / "sha256" / manifest_hash, "w") as mf:
        mf.write(manifest_json)
    return manifest_hash


def check_and_sanitize_flag_add_modelpack(add_modelpack: typing.Union[bool, None], ocilayout_indexes: Dict[str, OCIImageIndex], ocilayout_manifests: Dict[str, OCIImageManifest]) -> bool:
    """Check and sanitize the add_modelpack flag
    - check if the oci-layout contains an Index manifest for multi-arch, othewise fail: can't add a ModelPack manifest to a single-arch oci-layout
      this is because a single-arch is a single OCI Image Manifest and not introducing an Index; hence, can't add a ModelPack manifest to the non-existing Index which is tagged (`:latest`) by the oci-layout root index
    - sanitize the add_modelpack flag to avoid adding a ModelPack manifest if it's already present in the oci-layout
    """
    if add_modelpack:
        if len(ocilayout_indexes) == 0:
            raise ValueError("Can't add a ModelPack manifest to a single-arch oci-layout")
        logger.debug("checking if ModelPack manifest is alreadypresent in the oci-layout")
        for manifest_hash, manifest in ocilayout_manifests.items():
            if manifest.artifactType == modelpack_consts.ARTIFACTTYPEMODELMANIFEST:
                logger.warning("ModelPack manifest already present in the oci-layout as %s; will IGNORE add_modelpack=True", manifest_hash)
                return False
    return add_modelpack or False # leave unchanged unless of sanitization requirements above, if None just return False


def check_manifest(manifest: OCIImageManifest, config: OCIManifestConfig):
    """perform some sanity check on the manifests required for additional scenarios of usage
    """
    ch_count = len(list(x for x in config.history if not x.empty_layer)) if config.history else 0
    layers_count = len(manifest.layers)
    if layers_count != ch_count:
        raise ValueError(f"history lists {ch_count} non-empty layers, but there are {layers_count} layers in the image manifest")


def crawl_ocilayout_manifests(ocilayout: Path, ocilayout_indexes: Dict[str, OCIImageIndex], ocilayout_root_index: typing.Union[OCIImageIndex, None] = None) -> Dict[str, OCIImageManifest]:
    """crawl Manifests from referred OCI Index(es) and Manifests in the root index of the oci-layout
    """
    ocilayout_manifests: Dict[str, OCIImageManifest]  = {}
    for _, mi in ocilayout_indexes.items():
        for m in mi.manifests:
            logger.debug("Parsing manifest %s", m)
            if m.mediaType != MediaTypes.manifest:
                raise ValueError("Did not expect something else than Image Manifest in a Index")
            target_hash = m.digest.removeprefix("sha256:")
            logger.debug("target_hash %s", target_hash)
            manifest_path = ocilayout / "blobs" / "sha256" / target_hash
            with open(manifest_path, "r") as ip:
                ocilayout_manifests[target_hash] = OCIImageManifest.model_validate_json(ip.read())
    for m in ocilayout_root_index.manifests if ocilayout_root_index is not None else []:
        if m.mediaType == MediaTypes.manifest:
            target_hash = m.digest.removeprefix("sha256:")
            logger.debug("Lookup remainder manifest from ocilayout_root_index having target_hash %s", target_hash)
            manifest_path = ocilayout / "blobs" / "sha256" / target_hash
            with open(manifest_path, "r") as ip:
                ocilayout_manifests[target_hash] = OCIImageManifest.model_validate_json(ip.read())

    # filter out non-runnable OCI Images, like Vendor'd Attestations format, and log it out
    filtered: Dict[str, OCIImageManifest]  = {}
    for k, v in ocilayout_manifests.items():
        if v.layers[0].mediaType == "application/vnd.in-toto+json" or v.artifactType == "application/vnd.docker.attestation.manifest.v1+json":
            logger.info("skipping %s as it's an Attestation manifest", k) # not adding this to filtered list of manifests.
        else:
            filtered[k] = v

    return filtered


def write_empty_config_in_ocilayoyt(ocilayout: Path):
    """small utility to avoid limitation of skopeo that can't read inline empty config
    """
    blobs_path = ocilayout / "blobs" / "sha256"
    blobs_path.mkdir(parents=True, exist_ok=True)
    with open(blobs_path / "44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a", 'w') as f:
        f.write("{}")


def crawl_ocilayout_indexes(ocilayout: Path, ocilayout_root_index: OCIImageIndex) -> Dict[str, OCIImageIndex] :
    ocilayout_indexes: Dict[str, OCIImageIndex] = {}
    for m in ocilayout_root_index.manifests:
        if m.mediaType == MediaTypes.index:
            target_hash = m.digest.removeprefix("sha256:")
            index_path = ocilayout / "blobs" / "sha256" / target_hash
            with open(index_path, "r") as ip:
                ocilayout_indexes[target_hash] = OCIImageIndex.model_validate_json(ip.read())
    return ocilayout_indexes


def crawl_ocilayout_blobs_to_extract(ocilayout: Path, 
                                     output_path: Path,
                                     tar_filter_dir: str = "/models") -> List[str]:
    """
    Extract from OCI Image/ModelCar only the contents from a specific directory.

    Args:
        ocilayout: The directory containing the oci-layout of the OCI Image/ModelCar.
        output_path: The directory where to extract the ML model assets from the ModelCar to.
        tar_filter_dir: The subdirectory in the ModelCar to extract, defaults to `"/models"`.

    Returns:
        The list of extracted ML contents from the OCI Image/ModelCar.
    """
    extracted: List[str] = []
    tar_filter_dir= tar_filter_dir.lstrip("/")
    blobs_path = ocilayout / "blobs" / "sha256"
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    verify_ocilayout(ocilayout)
    ocilayout_root_index = read_ocilayout_root_index(ocilayout)
    if len(ocilayout_root_index.manifests) != 1:
        raise ValueError("TODO the root index has more than one manifest, expected single ModelCar")
    manifest0 = ocilayout_root_index.manifests[0]
    if manifest0.mediaType != MediaTypes.manifest:
        raise ValueError("Can only extract from ModelCar Image manifests")
    target_hash = manifest0.digest.removeprefix("sha256:")
    manifest_path = blobs_path / target_hash
    with open(manifest_path, "r") as ip:
        image_manifest = OCIImageManifest.model_validate_json(ip.read())
    for layer in image_manifest.layers:
        if layer.mediaType == MediaTypes.layer or layer.mediaType == MediaTypes.layer_gzip:
            target_hash = layer.digest.removeprefix("sha256:")
            manifest_path = blobs_path / target_hash
            with tarfile.open(manifest_path, "r:*") as tar:
                for member in tar.getmembers():
                    if member.isfile() and member.name.startswith(tar_filter_dir):
                        tar.extract(member, path=output_path)
                        extracted.append(member.name)
    return extracted
