"""Mock transport that runs workers locally in isolated temp directories.

Each MockConnection gets its own temp dir standing in for /workspace/.
push/pull copy files into/out of that dir. exec runs commands as local
subprocesses with /workspace paths rewritten to the temp dir.

Supports injecting failures: call conn.kill() to simulate a node dying.
"""

from __future__ import annotations

import asyncio
import os
import shutil
import tempfile


class MockConnection:
    def __init__(self, root_dir: str):
        self.root_dir = root_dir
        self._alive = True

    async def exec(self, cmd: str, timeout: float = 30.0) -> tuple[int, str]:
        if not self._alive:
            raise ConnectionError("node is dead")

        cmd = cmd.replace("/workspace", self.root_dir)

        proc = await asyncio.create_subprocess_shell(
            cmd,
            cwd=self.root_dir,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        try:
            stdout, stderr = await asyncio.wait_for(
                proc.communicate(), timeout=timeout
            )
            return proc.returncode, (stdout + stderr).decode()
        except asyncio.TimeoutError:
            proc.kill()
            return 1, "timeout"

    async def push(self, local_path: str, remote_path: str) -> None:
        if not self._alive:
            raise ConnectionError("node is dead")
        dest = remote_path.replace("/workspace", self.root_dir)
        os.makedirs(os.path.dirname(dest) or self.root_dir, exist_ok=True)
        shutil.copy2(local_path, dest)

    async def pull(self, remote_path: str, local_path: str) -> None:
        if not self._alive:
            raise ConnectionError("node is dead")
        src = remote_path.replace("/workspace", self.root_dir)
        if not os.path.exists(src):
            raise FileNotFoundError(f"remote file not found: {src}")
        os.makedirs(os.path.dirname(local_path) or ".", exist_ok=True)
        shutil.copy2(src, local_path)

    async def close(self) -> None:
        pass

    def kill(self):
        """Simulate node death."""
        self._alive = False

    def revive(self):
        """Bring node back (for testing reconnect scenarios)."""
        self._alive = True


class MockTransport:
    """Creates MockConnections backed by temp directories."""

    def __init__(self):
        self._dirs: list[str] = []

    async def connect(self, host: str, port: int) -> MockConnection:
        d = tempfile.mkdtemp(prefix=f"mock_node_{host}_{port}_")
        self._dirs.append(d)
        return MockConnection(d)

    def cleanup(self):
        for d in self._dirs:
            shutil.rmtree(d, ignore_errors=True)