# ========= 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. =========

import asyncio
import importlib.metadata
import json
import tarfile

import pytest
from pydantic import BaseModel, SecretStr

from mirage.commands.cli.specs import register_cli_spec, unregister_cli_spec
from mirage.commands.cli.types import CLIInvocation, CLISpec
from mirage.io import IOResult
from mirage.resource.disk import DiskResource
from mirage.resource.ram import RAMResource
from mirage.resource.s3 import S3Config, S3Resource
from mirage.resource.secrets import REDACTED_SECRET
from mirage.runtime.types import ScriptSource
from mirage.types import MountMode
from mirage.workspace import Workspace
from mirage.workspace.snapshot import to_state_dict
from mirage.workspace.snapshot.keys import CLIKey, ScriptKey, StateKey
from mirage.workspace.snapshot.utils import FORMAT_VERSION


def _load(*args, **kwargs):
    return asyncio.run(Workspace.load(*args, **kwargs))


def _seed(ws, mount: str = "/m") -> None:

    async def _do():
        await ws.execute(f"echo hello > {mount}/a.txt")
        await ws.execute(
            f"mkdir -p {mount}/sub && echo world > {mount}/sub/b.txt")

    asyncio.run(_do())


def _read(ws, path: str) -> str:

    async def _do():
        r = await ws.execute(f"cat {path}")
        return await r.stdout_str()

    return asyncio.run(_do())


# ── RAM round trip ──────────────────────────────────────────────────


def test_save_load_ram_round_trip(tmp_path):
    src = Workspace({"/m": (RAMResource(), MountMode.WRITE)},
                    mode=MountMode.WRITE)
    _seed(src)
    snap = tmp_path / "ram.tar"
    asyncio.run(src.snapshot(snap))
    assert snap.exists() and snap.stat().st_size > 0

    dst = _load(snap)
    assert _read(dst, "/m/a.txt") == "hello\n"
    assert _read(dst, "/m/sub/b.txt") == "world\n"


def test_history_survives_snapshot_round_trip(tmp_path):
    src = Workspace({"/m": (RAMResource(), MountMode.WRITE)},
                    mode=MountMode.WRITE)
    asyncio.run(src.execute("echo one"))
    asyncio.run(src.execute("echo two"))
    assert len(asyncio.run(src.history())) == 2
    snap = tmp_path / "history.tar"
    asyncio.run(src.snapshot(snap))

    dst = _load(snap)
    entries = asyncio.run(dst.history())
    assert [e["command"] for e in entries] == ["echo one", "echo two"]


def test_from_state_rebuilds_in_process_without_tar():
    src = Workspace({"/m": (RAMResource(), MountMode.WRITE)},
                    mode=MountMode.WRITE)
    _seed(src)

    dst = asyncio.run(Workspace.from_state(asyncio.run(to_state_dict(src))))
    assert _read(dst, "/m/a.txt") == "hello\n"
    assert _read(dst, "/m/sub/b.txt") == "world\n"


def test_save_load_ram_compressed_gz(tmp_path):
    src = Workspace({"/m": (RAMResource(), MountMode.WRITE)},
                    mode=MountMode.WRITE)
    _seed(src)
    snap = tmp_path / "ram.tar.gz"
    asyncio.run(src.snapshot(snap, compress="gz"))

    dst = _load(snap)
    assert _read(dst, "/m/a.txt") == "hello\n"


# ── Disk round trip ────────────────────────────────────────────────


def test_save_load_disk_round_trip(tmp_path):
    src_root = tmp_path / "src"
    src_root.mkdir()
    src = Workspace(
        {"/m": (DiskResource(root=str(src_root)), MountMode.WRITE)},
        mode=MountMode.WRITE)
    _seed(src)
    snap = tmp_path / "disk.tar"
    asyncio.run(src.snapshot(snap))

    dst = _load(snap)
    assert _read(dst, "/m/a.txt") == "hello\n"
    assert _read(dst, "/m/sub/b.txt") == "world\n"


def test_save_load_disk_with_override_root(tmp_path):
    src_root = tmp_path / "src"
    src_root.mkdir()
    src = Workspace(
        {"/m": (DiskResource(root=str(src_root)), MountMode.WRITE)},
        mode=MountMode.WRITE)
    _seed(src)
    snap = tmp_path / "disk.tar"
    asyncio.run(src.snapshot(snap))

    dst_root = tmp_path / "dst"
    dst_root.mkdir()
    dst = _load(snap, resources={"/m": DiskResource(root=str(dst_root))})
    assert (dst_root / "a.txt").read_bytes() == b"hello\n"
    assert (dst_root / "sub" / "b.txt").read_bytes() == b"world\n"
    assert _read(dst, "/m/a.txt") == "hello\n"


# ── redacted secret override enforcement ─────────────────────────────


def test_redacted_secret_missing_resource_raises(tmp_path):
    cfg = S3Config(bucket="b",
                   region="us-east-1",
                   aws_access_key_id="AKIA-LEAK",
                   aws_secret_access_key="SECRET-LEAK")
    src = Workspace({"/s3": (S3Resource(cfg), MountMode.WRITE)},
                    mode=MountMode.WRITE)
    snap = tmp_path / "s3.tar"
    asyncio.run(src.snapshot(snap))

    with pytest.raises(ValueError, match=r"resources="):
        _load(snap)


def test_redacted_secret_lists_all_missing(tmp_path):
    src = Workspace(
        {
            "/ram": (RAMResource(), MountMode.WRITE),
            "/s3a": (S3Resource(
                S3Config(bucket="a",
                         region="us-east-1",
                         aws_access_key_id="x",
                         aws_secret_access_key="x")), MountMode.WRITE),
            "/s3b": (S3Resource(
                S3Config(bucket="b",
                         region="us-east-1",
                         aws_access_key_id="x",
                         aws_secret_access_key="x")), MountMode.WRITE),
        },
        mode=MountMode.WRITE)
    snap = tmp_path / "two-s3.tar"
    asyncio.run(src.snapshot(snap))

    with pytest.raises(ValueError) as ei:
        _load(snap)
    msg = str(ei.value)
    assert "/s3a" in msg
    assert "/s3b" in msg


def test_s3_without_inline_secret_loads_without_override(tmp_path):
    cfg = S3Config(bucket="b", region="us-east-1", aws_profile="dev")
    src = Workspace({"/s3": (S3Resource(cfg), MountMode.WRITE)},
                    mode=MountMode.WRITE)
    snap = tmp_path / "s3-profile.tar"
    asyncio.run(src.snapshot(snap))

    dst = _load(snap)
    mount = dst._registry.mount_for("/s3/")
    assert isinstance(mount.resource, S3Resource)
    assert mount.resource.config.aws_profile == "dev"


# ── cred redaction in raw bytes ─────────────────────────────────────


def test_no_real_creds_in_tar_bytes(tmp_path):
    cfg = S3Config(bucket="b",
                   region="us-east-1",
                   aws_access_key_id="AKIA-OBVIOUS-LEAK",
                   aws_secret_access_key="SECRET-OBVIOUS-LEAK")
    src = Workspace({"/s3": (S3Resource(cfg), MountMode.WRITE)},
                    mode=MountMode.WRITE)
    snap = tmp_path / "s3.tar"
    asyncio.run(src.snapshot(snap))

    raw = snap.read_bytes()
    assert b"AKIA-OBVIOUS-LEAK" not in raw
    assert b"SECRET-OBVIOUS-LEAK" not in raw
    assert b"<REDACTED>" in raw


# ── manifest validity ──────────────────────────────────────────────


def test_manifest_is_valid_json(tmp_path):
    src = Workspace({"/m": (RAMResource(), MountMode.WRITE)},
                    mode=MountMode.WRITE)
    _seed(src)
    snap = tmp_path / "snap.tar"
    asyncio.run(src.snapshot(snap))

    with tarfile.open(snap, "r") as tar:
        f = tar.extractfile("manifest.json")
        manifest = json.loads(f.read().decode("utf-8"))
    assert manifest["version"] == FORMAT_VERSION
    assert "mounts" in manifest
    assert "cache" in manifest


def test_disk_files_extractable_from_tar(tmp_path):
    src_root = tmp_path / "src"
    src_root.mkdir()
    src = Workspace(
        {"/m": (DiskResource(root=str(src_root)), MountMode.WRITE)},
        mode=MountMode.WRITE)
    _seed(src)
    snap = tmp_path / "disk.tar"
    asyncio.run(src.snapshot(snap))

    extract = tmp_path / "extract"
    extract.mkdir()
    with tarfile.open(snap, "r") as tar:
        for member in tar.getmembers():
            if member.name.startswith("mounts/0/files/"):
                tar.extract(member, extract, filter="data")
    assert (extract / "mounts/0/files/a.txt").read_bytes() == b"hello\n"
    assert (extract / "mounts/0/files/sub/b.txt").read_bytes() == b"world\n"


# ── path-traversal defense ─────────────────────────────────────────


def test_load_rejects_path_traversal_in_blob_ref(tmp_path):
    snap = tmp_path / "bad.tar"
    manifest = {
        "version":
        1,
        "mirage_version":
        "0.1.0",
        "default_session_id":
        "default",
        "default_agent_id":
        "default",
        "current_agent_id":
        "default",
        "sessions": [],
        "history":
        None,
        "mounts": [{
            "index": 0,
            "prefix": "/m",
            "mode": "WRITE",
            "consistency": "LAZY",
            "resource_class": "mirage.resource.ram.RAMResource",
            "resource_state": {
                "type": "ram",
                "files": {
                    "/x": {
                        "__file": "../../etc/passwd"
                    }
                },
                "dirs": [],
                "modified": {},
            },
        }],
        "cache": {
            "limit": 0,
            "max_drain_bytes": None,
            "entries": []
        },
        "jobs": [],
    }
    with tarfile.open(snap, "w") as tar:
        data = json.dumps(manifest).encode("utf-8")
        info = tarfile.TarInfo(name="manifest.json")
        info.size = len(data)
        import io as _io
        tar.addfile(info, _io.BytesIO(data))

    with pytest.raises(ValueError, match="Unsafe blob path"):
        _load(snap)


# ── copy ───────────────────────────────────────────────────────────


def test_workspace_copy_independence_ram():
    src = Workspace({"/m": (RAMResource(), MountMode.WRITE)},
                    mode=MountMode.WRITE)
    asyncio.run(src.execute("echo hi > /m/a.txt"))

    cp = asyncio.run(src.copy())
    asyncio.run(cp.execute("echo bye > /m/a.txt"))

    assert _read(src, "/m/a.txt") == "hi\n"
    assert _read(cp, "/m/a.txt") == "bye\n"


def test_workspace_copy_preserves_max_drain_bytes():
    src = Workspace({"/m": (RAMResource(), MountMode.WRITE)},
                    mode=MountMode.WRITE)
    src.max_drain_bytes = 1234
    cp = asyncio.run(src.copy())
    assert cp.max_drain_bytes == 1234


# ── state dict shape ───────────────────────────────────────────────


def test_to_state_dict_shape():
    ws = Workspace({"/m": (RAMResource(), MountMode.WRITE)},
                   mode=MountMode.WRITE)
    state = asyncio.run(to_state_dict(ws))
    assert state["version"] == FORMAT_VERSION
    assert state["mirage_version"] == importlib.metadata.version("mirage-ai")
    assert state["mirage_version"] != "unknown"
    assert isinstance(state["mounts"], list)
    assert state["cache"]["entries"] == []
    assert state["jobs"] == []


def test_snapshot_round_trip_no_sync_policy(tmp_path):
    ws = Workspace({"/data": RAMResource()})
    target = tmp_path / "snap.tar"
    asyncio.run(ws.snapshot(str(target)))
    restored = _load(str(target))
    assert restored is not None


# ── filenames with spaces / unicode ───────────────────────────────


def test_ram_round_trip_filenames_with_spaces(tmp_path):
    src = Workspace({"/m": (RAMResource(), MountMode.WRITE)},
                    mode=MountMode.WRITE)
    src._registry.mount_for("/m/").resource._store.files["/my file.txt"] = (
        b"with spaces")
    src._registry.mount_for("/m/").resource._store.files[
        "/dir with space/data.txt"] = b"nested with space"
    src._registry.mount_for("/m/").resource._store.files["/数据.txt"] = (
        "你好".encode())

    snap = tmp_path / "spaces.tar"
    asyncio.run(src.snapshot(snap))
    dst = _load(snap)

    files = dst._registry.mount_for("/m/").resource._store.files
    assert files["/my file.txt"] == b"with spaces"
    assert files["/dir with space/data.txt"] == b"nested with space"
    assert files["/数据.txt"].decode() == "你好"


def test_disk_round_trip_filenames_with_spaces(tmp_path):
    src_root = tmp_path / "src"
    src_root.mkdir()
    (src_root / "my file.txt").write_bytes(b"hello space")
    (src_root / "dir with space").mkdir()
    (src_root / "dir with space" / "data.txt").write_bytes(b"deep space")
    (src_root / "数据.txt").write_bytes("你好".encode())

    src = Workspace(
        {"/m": (DiskResource(root=str(src_root)), MountMode.WRITE)},
        mode=MountMode.WRITE)
    snap = tmp_path / "disk-spaces.tar"
    asyncio.run(src.snapshot(snap))

    dst_root = tmp_path / "dst"
    dst_root.mkdir()
    _load(snap, resources={"/m": DiskResource(root=str(dst_root))})
    assert (dst_root / "my file.txt").read_bytes() == b"hello space"
    assert ((dst_root / "dir with space" /
             "data.txt").read_bytes() == b"deep space")
    assert (dst_root / "数据.txt").read_bytes().decode() == "你好"


def test_is_safe_blob_path_allows_spaces_and_unicode():
    from mirage.workspace.snapshot import is_safe_blob_path
    assert is_safe_blob_path("my file.txt")
    assert is_safe_blob_path("dir with space/data.txt")
    assert is_safe_blob_path("数据.txt")
    assert not is_safe_blob_path("../etc/passwd")
    assert not is_safe_blob_path("/abs/path")
    assert not is_safe_blob_path("")
    assert not is_safe_blob_path("foo/../bar")
    assert not is_safe_blob_path("foo\x00bar")


# ── Redis round trip ───────────────────────────────────────────────


@pytest.mark.skipif(not __import__("os").environ.get("REDIS_URL"),
                    reason="REDIS_URL not set")
def test_redis_round_trip_filenames_with_spaces(tmp_path):
    import os
    import uuid

    import redis as sync_redis

    from mirage.resource.redis import RedisResource
    redis_url = os.environ["REDIS_URL"]
    src_prefix = f"mirage:test:src:{uuid.uuid4().hex}:"
    dst_prefix = f"mirage:test:dst:{uuid.uuid4().hex}:"

    sc = sync_redis.Redis.from_url(redis_url)
    sc.set(f"{src_prefix}file:/my file.txt", b"hello space")
    sc.set(f"{src_prefix}file:/dir with space/data.txt", b"deep space")
    sc.sadd(f"{src_prefix}dir", "/dir with space")
    sc.close()

    src = Workspace(
        {
            "/m": (RedisResource(url=redis_url,
                                 key_prefix=src_prefix), MountMode.WRITE)
        },
        mode=MountMode.WRITE)
    snap = tmp_path / "redis-spaces.tar"
    asyncio.run(src.snapshot(snap))

    dst_resource = RedisResource(url=redis_url, key_prefix=dst_prefix)
    _load(snap, resources={"/m": dst_resource})

    sc = sync_redis.Redis.from_url(redis_url)
    try:
        assert sc.get(f"{dst_prefix}file:/my file.txt") == b"hello space"
        assert (sc.get(f"{dst_prefix}file:/dir with space/data.txt") ==
                b"deep space")
    finally:
        for prefix in (src_prefix, dst_prefix):
            for key in sc.scan_iter(f"{prefix}*"):
                sc.delete(key)
        sc.close()


class _CliCfg(BaseModel):
    token: SecretStr
    channel: str = "general"


async def _cli_echo(inv: CLIInvocation):
    return (f"tok={inv.config.token.get_secret_value()}\n".encode(),
            IOResult())


_CLI_SPEC = CLISpec(name="snapcli",
                    config_model=_CliCfg,
                    subcommands=(CLISpec(name="run", fn=_cli_echo), ))


@pytest.mark.asyncio
async def test_cli_registry_snapshots_with_redacted_config():
    register_cli_spec(_CLI_SPEC)
    try:
        ws = Workspace({"/data": RAMResource()}, mode=MountMode.WRITE)
        ws.register_cli("snapcli",
                        _CLI_SPEC,
                        config={
                            "token": "sek",
                            "channel": "eng"
                        })
        state = await to_state_dict(ws)
        entry = state[StateKey.CLIS][0]
        assert entry[CLIKey.NAME] == "snapcli"
        assert entry[CLIKey.SPEC] == "snapcli"
        assert entry[CLIKey.CONFIG] == {
            "token": REDACTED_SECRET,
            "channel": "eng"
        }

        with pytest.raises(ValueError, match="clis= must include"):
            await Workspace.from_state(state)

        ws2 = await Workspace.from_state(
            state, clis={"snapcli": {
                "token": "sek2",
                "channel": "eng"
            }})
        io = await ws2.execute("snapcli run")
        assert io.exit_code == 0
        assert io.stdout == b"tok=sek2\n"
        await ws.close()
        await ws2.close()
    finally:
        unregister_cli_spec("snapcli")


@pytest.mark.asyncio
async def test_copy_shares_live_cli_secrets():
    register_cli_spec(_CLI_SPEC)
    try:
        ws = Workspace({"/data": RAMResource()}, mode=MountMode.WRITE)
        ws.register_cli("snapcli", _CLI_SPEC, config={"token": "sek"})
        clone = await ws.copy()
        io = await clone.execute("snapcli run")
        assert io.exit_code == 0
        assert io.stdout == b"tok=sek\n"
        await ws.close()
        await clone.close()
    finally:
        unregister_cli_spec("snapcli")


@pytest.mark.asyncio
async def test_copy_carries_a_directly_installed_spec():
    # The spec is never named in the global registry: copy() must share
    # the live spec like a live resource, not resolve it by name.
    ws = Workspace({"/data": RAMResource()}, mode=MountMode.WRITE)
    ws.register_cli("snapcli", _CLI_SPEC, config={"token": "sek"})
    clone = await ws.copy()
    io = await clone.execute("snapcli run")
    assert io.exit_code == 0
    assert io.stdout == b"tok=sek\n"
    await ws.close()
    await clone.close()


class _NestedAuth(BaseModel):
    token: SecretStr


class _NestedCfg(BaseModel):
    auth: _NestedAuth


async def _nested_echo(inv: CLIInvocation):
    return (f"tok={inv.config.auth.token.get_secret_value()}\n".encode(),
            IOResult())


_NESTED_SPEC = CLISpec(name="nestcli",
                       config_model=_NestedCfg,
                       subcommands=(CLISpec(name="run", fn=_nested_echo), ))


@pytest.mark.asyncio
async def test_nested_cli_secrets_redact_and_demand_an_override():
    register_cli_spec(_NESTED_SPEC)
    try:
        ws = Workspace({"/data": RAMResource()}, mode=MountMode.WRITE)
        ws.register_cli("nestcli",
                        _NESTED_SPEC,
                        config={"auth": {
                            "token": "sek"
                        }})
        state = await to_state_dict(ws)
        entry = state[StateKey.CLIS][0]
        assert entry[CLIKey.CONFIG] == {"auth": {"token": REDACTED_SECRET}}
        with pytest.raises(ValueError, match="clis= must include"):
            await Workspace.from_state(state)
        clone = await ws.copy()
        io = await clone.execute("nestcli run")
        assert io.exit_code == 0
        assert io.stdout == b"tok=sek\n"
        await ws.close()
        await clone.close()
    finally:
        unregister_cli_spec("nestcli")


@pytest.mark.asyncio
async def test_script_cli_survives_a_tar_snapshot(tmp_path):
    # A script install resolves under no registry name, so the embedded
    # program rides in the snapshot and load rebuilds the spec from it.
    # The install also proves the manifest carries the clis key at all.
    ws = Workspace({"/data": RAMResource()}, mode=MountMode.WRITE)
    spec = CLISpec(name="pager",
                   script=ScriptSource(
                       "import os\n"
                       "print(argv[1], os.environ['MIRAGE_CLI_CONFIG'])"))
    ws.register_cli("pager", spec, config={"width": 80})
    path = tmp_path / "snap.tar"
    await ws.snapshot(str(path))
    with tarfile.open(path) as tf:
        manifest = json.loads(tf.extractfile("manifest.json").read())
    entry = manifest[StateKey.CLIS][0]
    assert entry[CLIKey.SPEC] == "pager"
    assert entry[CLIKey.SCRIPT][ScriptKey.LANGUAGE] == "python"
    assert entry[CLIKey.SCRIPT][ScriptKey.MODULE] is False

    restored = await Workspace.load(str(path))
    io = await restored.execute("pager b.txt")
    assert io.exit_code == 0
    assert io.stdout == b'b.txt {"width": 80}\n'
    await ws.close()
    await restored.close()


@pytest.mark.asyncio
async def test_script_cli_config_captures_verbatim_without_a_schema():
    # A script spec declares no config_model, so nothing declares which
    # keys are secret: the mapping is captured as-is rather than
    # crashing on model_dump or being guessed at.
    ws = Workspace({"/data": RAMResource()}, mode=MountMode.WRITE)
    spec = CLISpec(name="pager", script=ScriptSource("print('hi')"))
    ws.register_cli("pager", spec, config={"width": 80})
    state = await to_state_dict(ws)
    assert state[StateKey.CLIS][0][CLIKey.CONFIG] == {"width": 80}
    # copy() reveals secrets through the same helper, so it must not
    # crash on a mapping config either.
    clone = await ws.copy()
    io = await clone.execute("pager")
    assert io.exit_code == 0
    await ws.close()
    await clone.close()


@pytest.mark.asyncio
async def test_script_cli_runtime_pin_and_module_bit_round_trip():
    ws = Workspace({"/data": RAMResource()}, mode=MountMode.WRITE)
    spec = CLISpec(name="pager",
                   script=ScriptSource("x", language="js", module=True),
                   runtime="quickjs")
    ws.register_cli("pager", spec, config=None)
    state = await to_state_dict(ws)
    entry = state[StateKey.CLIS][0]
    assert entry[CLIKey.RUNTIME] == "quickjs"
    assert entry[CLIKey.SCRIPT][ScriptKey.MODULE] is True
    restored = await Workspace.from_state(state)
    install = restored.clis()["pager"]
    assert install.spec.runtime == "quickjs"
    assert install.spec.script.module is True
    assert install.spec.script.language == "js"
    await ws.close()
    await restored.close()
