# ========= 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 copy as _copy
import os
import shutil
import tempfile
from pathlib import Path

from mirage import MountMode, Workspace
from mirage.resource.disk import DiskResource

REPO_ROOT = Path(__file__).resolve().parents[3]
DATA_DIR = REPO_ROOT / "data"

tmp = tempfile.mkdtemp()
shutil.copytree(DATA_DIR, Path(tmp) / "files", dirs_exist_ok=True)

resource = DiskResource(root=tmp + "/files")


async def main() -> None:
    ws = Workspace({"/data/": resource}, mode=MountMode.READ)

    print("=== ls /data/ ===")
    result = await ws.execute("ls /data/")
    print(await result.stdout_str())

    print("=== cat /data/example.json ===")
    result = await ws.execute("cat /data/example.json")
    print(await result.stdout_str())

    print("=== head -n 3 /data/example.jsonl ===")
    result = await ws.execute("head -n 3 /data/example.jsonl")
    print(await result.stdout_str())

    print("=== tail -n 2 /data/example.jsonl ===")
    result = await ws.execute("tail -n 2 /data/example.jsonl")
    print(await result.stdout_str())

    print("=== wc /data/example.json ===")
    result = await ws.execute("wc /data/example.json")
    print(await result.stdout_str())

    print("=== stat /data/example.json ===")
    result = await ws.execute("stat /data/example.json")
    print(await result.stdout_str())

    print("=== tree /data/ ===")
    result = await ws.execute("tree /data/")
    print(await result.stdout_str())

    print("=== find /data/ -name '*.json' ===")
    result = await ws.execute("find /data/ -name '*.json'")
    print(await result.stdout_str())

    print("=== grep example /data/example.json ===")
    result = await ws.execute("grep example /data/example.json")
    print(await result.stdout_str())

    print("=== rg example /data/example.json ===")
    result = await ws.execute("rg example /data/example.json")
    print(await result.stdout_str())

    print("=== du /data/ ===")
    result = await ws.execute("du /data/")
    print(await result.stdout_str())

    print("=== jq .company /data/example.json ===")
    result = await ws.execute('jq ".company" /data/example.json')
    print(await result.stdout_str())

    print("=== basename /data/example.json ===")
    result = await ws.execute("basename /data/example.json")
    print(await result.stdout_str())

    print("=== dirname /data/example.json ===")
    result = await ws.execute("dirname /data/example.json")
    print(await result.stdout_str())

    print("=== sort /data/example.jsonl ===")
    result = await ws.execute("sort /data/example.jsonl")
    print(await result.stdout_str())

    print("=== cp /data/example.json /data/copy.json ===")
    await ws.execute("cp /data/example.json /data/copy.json")
    result = await ws.execute("ls /data/")
    print(await result.stdout_str())

    print("=== mv /data/copy.json /data/renamed.json ===")
    await ws.execute("mv /data/copy.json /data/renamed.json")
    result = await ws.execute("ls /data/")
    print(await result.stdout_str())

    print("=== rm /data/renamed.json ===")
    await ws.execute("rm /data/renamed.json")
    result = await ws.execute("ls /data/")
    print(await result.stdout_str())

    # ── persistence: save / load / copy / deepcopy ──────────────────
    # Disk has needs_override=False — full file tree is in the snapshot.
    # Default load behavior creates a fresh tmpdir. Caller can override
    # by supplying DiskResource(root=...) in resources={...}.
    print("\n=== PERSISTENCE ===\n")
    with tempfile.NamedTemporaryFile(suffix=".tar", delete=False) as f:
        snap = f.name
    custom_root = tempfile.mkdtemp(prefix="mirage-disk-restore-")
    try:
        ws.snapshot(snap)
        print(f"  saved → {snap} ({os.path.getsize(snap)} bytes)")

        # Load with default fresh tmpdir
        loaded_default = Workspace.load(snap)
        r = await loaded_default.execute("ls /data/")
        print(f"  loaded (default tmpdir) ls: "
              f"{(await r.stdout_str()).strip()[:80]}…")

        # Load with caller-supplied root — files written into custom_root
        loaded_custom = Workspace.load(
            snap, resources={"/data": DiskResource(root=custom_root)})
        r = await loaded_custom.execute("ls /data/")
        print(f"  loaded (root={custom_root[:40]}…) ls: "
              f"{(await r.stdout_str()).strip()[:80]}…")
        print(f"  custom_root contents: {sorted(os.listdir(custom_root))[:5]}")

        cp = ws.copy()
        print(f"  copy() mounts: {[m.prefix for m in cp.mounts()]}")

        deep = _copy.deepcopy(ws)
        print(f"  deepcopy() mounts: {[m.prefix for m in deep.mounts()]}")

        try:
            _copy.copy(ws)
            print("  ✗ shallow copy should have raised")
        except NotImplementedError as e:
            print(f"  ✓ shallow copy raises: {str(e)[:60]}…")
    finally:
        os.unlink(snap)
        shutil.rmtree(custom_root, ignore_errors=True)


if __name__ == "__main__":
    asyncio.run(main())
