import base64
import io
import json
import mimetypes
import os
import zipfile
from typing import Any


def unpack_poma_archive(
    poma_archive_data: bytes | None = None,
    poma_archive_path: str | os.PathLike[str] | None = None,
) -> dict[str, Any]:
    """
    Extract chunks, chunksets, image sources, and raw asset files from a
    POMA archive.

    Args:
        poma_archive_data: The POMA archive as bytes.
        poma_archive_path: Path to the POMA archive file.
    Returns:
        dict with keys ``chunks``, ``chunksets``, ``image_sources``
        (raw dict from image_sources.json, values may be base64 data URIs,
        URLs, or None), and ``asset_files`` (dict mapping asset paths to
        raw bytes from the ``assets/`` folder).
    """
    if poma_archive_path:
        zip_file = zipfile.ZipFile(poma_archive_path, "r")
    elif poma_archive_data:
        zip_file = zipfile.ZipFile(io.BytesIO(poma_archive_data), "r")
    else:
        raise ValueError(
            "Either poma_archive_data or poma_archive_path must be provided."
        )

    with zip_file as zf:
        chunks = zf.read("chunks.json")
        chunksets = zf.read("chunksets.json")

        image_sources: dict[str, Any] = {}
        if "image_sources.json" in zf.namelist():
            image_sources = json.loads(zf.read("image_sources.json"))

        asset_files: dict[str, bytes] = {}
        for name in zf.namelist():
            if name.startswith("assets/") and not name.endswith("/"):
                asset_files[name] = zf.read(name)

    if not chunks or not chunksets:
        raise KeyError("Result must contain 'chunks' and 'chunksets' keys.")
    return {
        "chunks": json.loads(chunks),
        "chunksets": json.loads(chunksets),
        "image_sources": image_sources,
        "asset_files": asset_files,
    }


def _resolve_images(
    image_sources: dict[str, Any],
    asset_files: dict[str, bytes],
) -> dict[str, str]:
    """
    Resolve image sources to base64 data URIs.

    For each entry in *image_sources*:
    - If the value is already a ``data:`` URI, keep it as-is.
    - If it is ``None`` or empty, try to find a matching file in
      *asset_files* and convert it to a base64 data URI.
    - Otherwise (e.g. a URL or unknown value), try to find a matching
      asset file as a fallback.

    Images that cannot be resolved are silently skipped.

    Returns:
        dict mapping image IDs to base64 data URI strings.
    """
    images: dict[str, str] = {}

    for image_id, value in image_sources.items():
        if isinstance(value, str) and value.startswith("data:"):
            images[image_id] = value
            continue

        # Try to find a matching asset file
        data_uri = _asset_to_data_uri(image_id, asset_files)
        if data_uri:
            images[image_id] = data_uri

    return images


def _asset_to_data_uri(
    image_id: str,
    asset_files: dict[str, bytes],
) -> str | None:
    """Find an asset file matching *image_id* and convert to a data URI."""
    for asset_path, asset_bytes in asset_files.items():
        # asset_path is e.g. "assets/image_00001.png"
        filename = asset_path.rsplit("/", 1)[-1]
        stem = filename.rsplit(".", 1)[0] if "." in filename else filename
        if stem == image_id:
            mime = mimetypes.guess_type(filename)[0] or "application/octet-stream"
            b64 = base64.b64encode(asset_bytes).decode("ascii")
            return f"data:{mime};base64,{b64}"
    return None


# Backwards-compatible alias
extract_chunks_and_chunksets_from_poma_archive = unpack_poma_archive
