from __future__ import annotations

from dataclasses import dataclass, field
import shlex
from typing import Mapping

import pytest
from shellctl.shared import DeleteJobResponse, JobResult, JobStatusName, JobStatusView

from dify_agent.runtime_backend import (
    BindingCreateError,
    BindingDestroyError,
    ExecutionBindingCreateSpec,
    ExecutionBindingDestroySpec,
    HomeSnapshotCreateSpec,
)
from dify_agent.runtime_backend.local import LocalExecutionBindingBackend, LocalHomeSnapshotBackend


@dataclass(slots=True)
class _RunCall:
    commands: tuple[tuple[str, ...], ...]
    cwd: str | None
    env: Mapping[str, str] | None


@dataclass(slots=True)
class _Client:
    runs: list[_RunCall]
    closed: bool = False
    exit_code: int = 0
    output: str = ""
    close_error: Exception | None = None
    exit_codes: list[int] = field(default_factory=list)

    async def run(
        self,
        script: str,
        *,
        cwd: str | None = None,
        env: Mapping[str, str] | None = None,
        timeout: float = 10.0,
    ) -> JobResult:
        del timeout
        commands = tuple(
            tuple(shlex.split(line)) for line in script.splitlines() if line.strip() and line.strip() != "set -eu"
        )
        self.runs.append(_RunCall(commands=commands, cwd=cwd, env=env))
        return JobResult(
            job_id=f"job-{len(self.runs)}",
            status=JobStatusName.EXITED,
            done=True,
            exit_code=self.exit_codes.pop(0) if self.exit_codes else self.exit_code,
            output_path="/tmp/output.log",
            output=self.output,
            offset=0,
            truncated=False,
        )

    async def wait(self, job_id: str, *, offset: int, timeout: float = 10.0) -> JobResult:
        raise AssertionError((job_id, offset, timeout))

    async def input(self, job_id: str, text: str, *, offset: int, timeout: float = 10.0) -> JobResult:
        raise AssertionError((job_id, text, offset, timeout))

    async def tail(self, job_id: str) -> JobResult:
        raise AssertionError(job_id)

    async def terminate(self, job_id: str, grace_seconds: float = 2.0) -> JobStatusView:
        raise AssertionError((job_id, grace_seconds))

    async def delete(
        self,
        job_id: str,
        *,
        force: bool = False,
        grace_seconds: float | None = None,
    ) -> DeleteJobResponse:
        del force, grace_seconds
        return DeleteJobResponse(job_id=job_id)

    async def close(self) -> None:
        self.closed = True
        if self.close_error is not None:
            raise self.close_error


@dataclass(slots=True)
class _Factory:
    clients: list[_Client] = field(default_factory=list)
    runs: list[_RunCall] = field(default_factory=list)

    def __call__(self) -> _Client:
        client = _Client(runs=self.runs)
        self.clients.append(client)
        return client

    @property
    def commands(self) -> tuple[tuple[str, ...], ...]:
        return tuple(command for run in self.runs for command in run.commands)


@dataclass(slots=True)
class _FailingFactory:
    clients: list[_Client] = field(default_factory=list)
    runs: list[_RunCall] = field(default_factory=list)

    def __call__(self) -> _Client:
        client = _Client(
            runs=self.runs,
            exit_code=1,
            output="primary shellctl failure",
            close_error=RuntimeError("secondary close failure"),
        )
        self.clients.append(client)
        return client


@dataclass(slots=True)
class _FailThenSucceedFactory:
    clients: list[_Client] = field(default_factory=list)
    runs: list[_RunCall] = field(default_factory=list)

    def __call__(self) -> _Client:
        client = _Client(
            runs=self.runs,
            output="primary shellctl failure",
            exit_codes=[1, 0],
        )
        self.clients.append(client)
        return client

    @property
    def commands(self) -> tuple[tuple[str, ...], ...]:
        return tuple(command for run in self.runs for command in run.commands)


@pytest.mark.anyio
async def test_local_binding_create_materializes_home_and_new_workspace() -> None:
    factory = _Factory()
    backend = LocalExecutionBindingBackend(
        endpoint="http://shellctl",
        auth_token="",
        materialized_home_root="/homes",
        workspace_root="/workspaces",
        snapshot_root="/snapshots",
        client_factory=factory,  # pyright: ignore[reportArgumentType]
    )

    allocation = await backend.create_binding(
        ExecutionBindingCreateSpec(
            tenant_id="tenant-1",
            agent_id="agent-1",
            binding_id="binding-1",
            workspace_id="workspace-1",
            existing_workspace_ref=None,
            home_snapshot_ref="home-home-1",
        )
    )

    assert allocation.binding_ref == "binding-1:workspace-1"
    assert allocation.workspace_ref == "workspace-1"
    assert factory.commands[0] == ("test", "-d", "/snapshots/home-home-1")
    assert ("mkdir", "-p", "/workspaces/workspace-1") in factory.commands
    assert ("mkdir", "-p", "/homes/binding-1") in factory.commands
    assert ("cp", "-a", "/snapshots/home-home-1/.", "/homes/binding-1/") in factory.commands
    assert ("chmod", "700", "/homes/binding-1", "/workspaces/workspace-1") in factory.commands


@pytest.mark.anyio
async def test_local_binding_create_uses_empty_default_home_without_snapshot_access() -> None:
    factory = _Factory()
    backend = LocalExecutionBindingBackend(
        endpoint="http://shellctl",
        auth_token="",
        materialized_home_root="/homes",
        workspace_root="/workspaces",
        snapshot_root="/snapshots",
        client_factory=factory,  # pyright: ignore[reportArgumentType]
    )

    allocation = await backend.create_binding(
        ExecutionBindingCreateSpec(
            tenant_id="tenant-1",
            agent_id="agent-1",
            binding_id="binding-1",
            workspace_id="workspace-1",
            existing_workspace_ref=None,
            home_snapshot_ref=None,
        )
    )

    assert allocation.binding_ref == "binding-1:workspace-1"
    assert ("mkdir", "-p", "/homes/binding-1") in factory.commands
    assert all("/snapshots" not in part for command in factory.commands for part in command)


@pytest.mark.anyio
async def test_local_binding_create_failure_removes_partial_home_and_workspace() -> None:
    factory = _FailThenSucceedFactory()
    backend = LocalExecutionBindingBackend(
        endpoint="http://shellctl",
        auth_token="",
        materialized_home_root="/homes",
        workspace_root="/workspaces",
        snapshot_root="/snapshots",
        client_factory=factory,  # pyright: ignore[reportArgumentType]
    )

    with pytest.raises(BindingCreateError, match="primary shellctl failure"):
        await backend.create_binding(
            ExecutionBindingCreateSpec(
                tenant_id="tenant-1",
                agent_id="agent-1",
                binding_id="binding-1",
                workspace_id="workspace-1",
                existing_workspace_ref=None,
                home_snapshot_ref="home-home-1",
            )
        )

    assert ("rm", "-rf", "--", "/homes/binding-1", "/workspaces/workspace-1") in factory.commands


@pytest.mark.anyio
async def test_local_binding_acquire_scopes_commands_to_materialized_home_and_workspace() -> None:
    factory = _Factory()
    backend = LocalExecutionBindingBackend(
        endpoint="http://shellctl",
        auth_token="",
        materialized_home_root="/homes",
        workspace_root="/workspaces",
        snapshot_root="/snapshots",
        client_factory=factory,  # pyright: ignore[reportArgumentType]
    )

    lease = await backend.acquire("binding-1:workspace-1")

    assert lease.layout.home_dir == "/homes/binding-1"
    assert lease.layout.workspace_dir == "/workspaces/workspace-1"
    assert ("test", "-d", "/homes/binding-1") in factory.commands
    assert ("test", "-d", "/workspaces/workspace-1") in factory.commands

    await lease.commands.run("pwd", cwd=None, env={"HOME": "/homes/other"}, timeout=10.0)
    pwd_run = next(run for run in factory.runs if run.commands == (("pwd",),))
    assert pwd_run.cwd == "/workspaces/workspace-1"
    assert pwd_run.env == {"HOME": "/homes/binding-1"}
    with pytest.raises(ValueError, match="outside this RuntimeLease"):
        await lease.commands.run("cat secret", cwd="/homes/other", timeout=10.0)
    await backend.release(lease)


@pytest.mark.anyio
async def test_local_snapshot_checkpoint_copies_only_materialized_home() -> None:
    factory = _Factory()
    snapshots = LocalHomeSnapshotBackend(
        endpoint="http://shellctl",
        auth_token="",
        snapshot_root="/snapshots",
        client_factory=factory,  # pyright: ignore[reportArgumentType]
    )
    bindings = LocalExecutionBindingBackend(
        endpoint="http://shellctl",
        auth_token="",
        materialized_home_root="/homes",
        workspace_root="/workspaces",
        snapshot_root="/snapshots",
        client_factory=factory,  # pyright: ignore[reportArgumentType]
    )
    lease = await bindings.acquire("binding-1:workspace-1")

    snapshot_ref = await snapshots.create_from_runtime(
        spec=HomeSnapshotCreateSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-2"),
        source=lease,
    )
    await bindings.release(lease)

    assert snapshot_ref == "home-home-2"
    assert ("test", "-d", "/homes/binding-1") in factory.commands
    assert ("mkdir", "-p", "/snapshots/home-home-2") in factory.commands
    assert ("cp", "-a", "/homes/binding-1/.", "/snapshots/home-home-2/") in factory.commands
    assert ("cp", "-a", "/workspaces/workspace-1/.", "/snapshots/home-home-2/") not in factory.commands


@pytest.mark.anyio
async def test_local_binding_destroy_removes_home_and_requested_workspace() -> None:
    factory = _Factory()
    backend = LocalExecutionBindingBackend(
        endpoint="http://shellctl",
        auth_token="",
        materialized_home_root="/homes",
        workspace_root="/workspaces",
        snapshot_root="/snapshots",
        client_factory=factory,  # pyright: ignore[reportArgumentType]
    )

    await backend.destroy_binding(
        ExecutionBindingDestroySpec(
            binding_ref="binding-1:workspace-1",
            workspace_ref="workspace-1",
            destroy_workspace=True,
        )
    )

    assert ("rm", "-rf", "--", "/homes/binding-1", "/workspaces/workspace-1") in factory.commands


@pytest.mark.anyio
async def test_local_snapshot_delete_removes_snapshot_directory() -> None:
    factory = _Factory()
    backend = LocalHomeSnapshotBackend(
        endpoint="http://shellctl",
        auth_token="",
        snapshot_root="/snapshots",
        client_factory=factory,  # pyright: ignore[reportArgumentType]
    )

    await backend.delete("home-home-2")

    assert ("rm", "-rf", "--", "/snapshots/home-home-2") in factory.commands


@pytest.mark.anyio
async def test_local_backend_materializes_same_agent_twice_in_one_workspace() -> None:
    factory = _Factory()
    backend = LocalExecutionBindingBackend(
        endpoint="http://shellctl",
        auth_token="",
        client_factory=factory,  # pyright: ignore[reportArgumentType]
    )

    first = await backend.create_binding(
        ExecutionBindingCreateSpec(
            tenant_id="tenant-1",
            agent_id="agent-1",
            binding_id="binding-1",
            workspace_id="workspace-1",
            existing_workspace_ref=None,
            home_snapshot_ref="home-home-1",
        )
    )
    second = await backend.create_binding(
        ExecutionBindingCreateSpec(
            tenant_id="tenant-1",
            agent_id="agent-1",
            binding_id="binding-2",
            workspace_id="workspace-1",
            existing_workspace_ref=first.workspace_ref,
            home_snapshot_ref="home-home-1",
        )
    )

    assert first.binding_ref == "binding-1:workspace-1"
    assert second.binding_ref == "binding-2:workspace-1"
    assert first.workspace_ref == second.workspace_ref == "workspace-1"
    first_lease = await backend.acquire(first.binding_ref)
    second_lease = await backend.acquire(second.binding_ref)
    assert first_lease.layout.home_dir != second_lease.layout.home_dir
    assert first_lease.layout.workspace_dir == second_lease.layout.workspace_dir
    await backend.release(first_lease)
    await backend.release(second_lease)

    await backend.destroy_binding(
        ExecutionBindingDestroySpec(
            binding_ref=first.binding_ref,
            destroy_workspace=False,
        )
    )
    surviving_lease = await backend.acquire(second.binding_ref)
    assert surviving_lease.layout.home_dir == "/home/dify/.dify-agent-materialized-homes/binding-2"
    assert surviving_lease.layout.workspace_dir == "/home/dify/.dify-agent-workspaces/workspace-1"
    await backend.release(surviving_lease)

    workspace_dir = "/home/dify/.dify-agent-workspaces/workspace-1"
    assert ("test", "-d", workspace_dir) in factory.commands
    assert factory.commands.count(("mkdir", "-p", workspace_dir)) == 1
    assert (
        "rm",
        "-rf",
        "--",
        "/home/dify/.dify-agent-materialized-homes/binding-1",
    ) in factory.commands
    assert (
        "rm",
        "-rf",
        "--",
        "/home/dify/.dify-agent-materialized-homes/binding-1",
        workspace_dir,
    ) not in factory.commands
    assert (
        "cp",
        "-a",
        "/home/dify/.dify-agent-home-snapshots/home-home-1/.",
        "/home/dify/.dify-agent-materialized-homes/binding-1/",
    ) in factory.commands
    assert (
        "cp",
        "-a",
        "/home/dify/.dify-agent-home-snapshots/home-home-1/.",
        "/home/dify/.dify-agent-materialized-homes/binding-2/",
    ) in factory.commands


@pytest.mark.anyio
async def test_local_snapshot_delete_preserves_shellctl_error_when_close_fails() -> None:
    factory = _FailingFactory()
    backend = LocalHomeSnapshotBackend(
        endpoint="http://shellctl",
        auth_token="",
        client_factory=factory,  # pyright: ignore[reportArgumentType]
    )

    with pytest.raises(BindingDestroyError, match="primary shellctl failure"):
        await backend.delete("home-1")

    assert factory.clients and all(client.closed for client in factory.clients)


@pytest.mark.anyio
async def test_local_binding_destroy_preserves_shellctl_error_when_close_fails() -> None:
    factory = _FailingFactory()
    backend = LocalExecutionBindingBackend(
        endpoint="http://shellctl",
        auth_token="",
        client_factory=factory,  # pyright: ignore[reportArgumentType]
    )

    with pytest.raises(BindingDestroyError, match="primary shellctl failure"):
        await backend.destroy_binding(
            ExecutionBindingDestroySpec(
                binding_ref="binding-1:workspace-1",
                destroy_workspace=False,
            )
        )

    assert factory.clients and all(client.closed for client in factory.clients)