# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========

from typing import Any

from mirage.types import ResourceName
from mirage.workspace.snapshot.keys import (CacheKey, JobKey, MountKey,
                                            ResourceStateKey, StateKey)
from mirage.workspace.snapshot.utils import BLOB_REF_KEY, is_safe_blob_path


class _BlobAllocator:
    """Mints unique blob filenames per category and collects their bytes.

    `alloc("_cache") -> "_cache/0.bin"` on first call, "_cache/1.bin"
    next, and so on. The category prefix is just for uniqueness; the
    real tar path is decided by the caller and stored in `blobs`.
    """

    def __init__(self) -> None:
        self.blobs: dict[str, bytes] = {}
        self._counters: dict[str, int] = {}

    def alloc(self, category: str) -> str:
        i = self._counters.get(category, 0)
        self._counters[category] = i + 1
        return f"{category}/{i}.bin"


def split_manifest_and_blobs(
        state: dict[str, Any]) -> tuple[dict[str, Any], dict[str, bytes]]:
    """The state as a JSON-safe manifest plus the bytes it referenced.

    The manifest IS the state, minus the four keys that hold bytes:
    those are rewritten with blob references and everything else rides
    through untouched. It is built that way round on purpose. An
    allowlist of keys silently dropped whatever it forgot from every
    tar snapshot while leaving in-memory snapshots (which never pass
    through here) green, which is exactly how installed CLIs went
    missing.

    Args:
        state (dict[str, Any]): the workspace state dict.
    """
    a = _BlobAllocator()
    manifest: dict[str, Any] = dict(state)
    cache = state.get(StateKey.CACHE) or {}

    manifest[StateKey.HISTORY] = _history_to_manifest(
        state.get(StateKey.HISTORY), a)
    manifest[StateKey.MOUNTS] = [
        _mount_to_manifest(m, a) for m in state.get(StateKey.MOUNTS) or []
    ]
    manifest[StateKey.CACHE] = {
        **cache,
        CacheKey.ENTRIES: [
            _cache_entry_to_manifest(e, a)
            for e in cache.get(CacheKey.ENTRIES) or []
        ],
    }
    manifest[StateKey.JOBS] = [
        _job_to_manifest(j, a) for j in state.get(StateKey.JOBS) or []
    ]

    return manifest, a.blobs


def _cache_entry_to_manifest(entry: dict[str, Any],
                             a: _BlobAllocator) -> dict[str, Any]:
    """One cache entry with its data stashed as a blob reference.

    Args:
        entry (dict[str, Any]): the captured cache entry.
        a (_BlobAllocator): the run's blob collector.
    """
    e = dict(entry)
    if isinstance(e.get(CacheKey.DATA), bytes):
        tar_path = "cache/blobs/" + a.alloc("_cache")
        a.blobs[tar_path] = e[CacheKey.DATA]
        e[CacheKey.DATA] = {BLOB_REF_KEY: tar_path}
    return e


def _job_to_manifest(job: dict[str, Any], a: _BlobAllocator) -> dict[str, Any]:
    """One job with its streams stashed as blob references.

    Args:
        job (dict[str, Any]): the captured job.
        a (_BlobAllocator): the run's blob collector.
    """
    j = dict(job)
    for f in (JobKey.STDOUT, JobKey.STDERR):
        data = j.get(f)
        if isinstance(data, bytes) and data:
            tar_path = "jobs/blobs/" + a.alloc("_jobs")
            a.blobs[tar_path] = data
            j[f] = {BLOB_REF_KEY: tar_path}
        elif isinstance(data, bytes):
            j[f] = ""  # empty bytes -> empty string (JSON-safe)
    return j


def _mount_to_manifest(mount: dict[str, Any],
                       a: _BlobAllocator) -> dict[str, Any]:
    idx = mount[MountKey.INDEX]
    ps = dict(mount[MountKey.RESOURCE_STATE])
    ptype = ps.get(ResourceStateKey.TYPE, "")
    files = ps.get(ResourceStateKey.FILES, {})
    if ptype == ResourceName.RAM:
        ps[ResourceStateKey.FILES] = _stash_blobs(
            files, a, f"_ram{idx}", tar_dir=f"mounts/{idx}/files")
    elif ptype == ResourceName.DISK:
        # tree-preserving: real files at their relative paths
        new_files: dict[str, dict[str, Any]] = {}
        for rel, data in files.items():
            tar_path = f"mounts/{idx}/files/{rel}"
            a.blobs[tar_path] = data
            new_files[rel] = {BLOB_REF_KEY: tar_path}
        ps[ResourceStateKey.FILES] = new_files
    elif ptype == ResourceName.REDIS:
        ps[ResourceStateKey.FILES] = _stash_blobs(files,
                                                  a,
                                                  f"_redis{idx}",
                                                  tar_dir=f"mounts/{idx}/data")
    return {
        **{
            k: v
            for k, v in mount.items() if k != MountKey.RESOURCE_STATE
        }, MountKey.RESOURCE_STATE: ps
    }


def _stash_blobs(files: dict[str, Any], a: _BlobAllocator, category: str,
                 tar_dir: str) -> dict[str, Any]:
    """Replace each {key: bytes} with {key: {__file: tar-path}}."""
    out: dict[str, dict[str, Any]] = {}
    for k, data in files.items():
        slot = a.alloc(category).split("/")[-1]
        tar_path = f"{tar_dir}/{slot}"
        a.blobs[tar_path] = data
        out[k] = {BLOB_REF_KEY: tar_path}
    return out


def _history_to_manifest(records, a: _BlobAllocator):
    if records is None:
        return None
    out = []
    for r in records:
        rd = dict(r)
        for f in ("stdout", "stdin", "stderr"):
            data = rd.get(f)
            if isinstance(data, bytes) and data:
                tar_path = "history/blobs/" + a.alloc("_history")
                a.blobs[tar_path] = data
                rd[f] = {BLOB_REF_KEY: tar_path}
            elif isinstance(data, bytes):
                rd[f] = ""  # empty bytes -> empty string
        if "tree" in rd:
            rd["tree"] = _node_to_manifest(rd["tree"], a)
        out.append(rd)
    return out


def _node_to_manifest(node, a: _BlobAllocator):
    if not isinstance(node, dict):
        return node
    out = dict(node)
    data = out.get("stderr")
    if isinstance(data, bytes) and data:
        tar_path = "history/blobs/" + a.alloc("_history")
        a.blobs[tar_path] = data
        out["stderr"] = {BLOB_REF_KEY: tar_path}
    elif isinstance(data, bytes):
        out["stderr"] = ""  # empty bytes -> empty string
    if "children" in out:
        out["children"] = [_node_to_manifest(c, a) for c in out["children"]]
    return out


def resolve_manifest(manifest: dict[str, Any], blob_reader) -> dict[str, Any]:
    return _resolve(manifest, blob_reader)


def _resolve(node, blob_reader):
    if isinstance(node, dict):
        if set(node.keys()) == {BLOB_REF_KEY}:
            path = node[BLOB_REF_KEY]
            if not is_safe_blob_path(path):
                raise ValueError(f"Unsafe blob path in manifest: {path!r}")
            return blob_reader(path)
        return {k: _resolve(v, blob_reader) for k, v in node.items()}
    if isinstance(node, list):
        return [_resolve(v, blob_reader) for v in node]
    return node
