import os
import types
import pytest
from typing import Optional

from pentest_mcp_server.server import PentestMCPServer


class DummySFTP:
	def __init__(self):
		self.operations = []
		self.closed = False

	async def put(self, local_path, remote_path):
		self.operations.append(("put", local_path, remote_path))

	async def get(self, remote_path, local_path):
		os.makedirs(os.path.dirname(local_path) or ".", exist_ok=True)
		with open(local_path, "wb") as f:
			f.write(b"data")
		self.operations.append(("get", remote_path, local_path))

	async def exit(self):
		self.closed = True


class DummySSHManager:
	def __init__(self, sftp: DummySFTP):
		self._sftp = sftp
		self.run_commands = []

	async def ensure_connected(self) -> bool:
		return True

	async def start_sftp_client(self):
		return self._sftp

	async def run_command(self, cmd: str, timeout: Optional[int] = None):
		self.run_commands.append(cmd)
		return types.SimpleNamespace(exit_status=0, stdout="", stderr="")


@pytest.mark.asyncio
async def test_upload_success(tmp_path):
	server = PentestMCPServer()
	sftp = DummySFTP()
	server.ssh_manager = DummySSHManager(sftp)
	server.tmux_manager = object()

	async def ok_init():
		return True

	server._ensure_initialized = ok_init  # type: ignore[attr-defined]

	local_file = tmp_path / "sample.txt"
	local_file.write_text("hello")

	result = await server._handle_upload_file({
		"local_path": str(local_file),
		"remote_path": "/tmp/sample.txt"
	})

	assert result["status"] == "success"
	assert ("put", str(local_file), "/tmp/sample.txt") in sftp.operations
	assert sftp.closed is True


@pytest.mark.asyncio
async def test_upload_missing_local():
	server = PentestMCPServer()
	sftp = DummySFTP()
	server.ssh_manager = DummySSHManager(sftp)
	server.tmux_manager = object()

	async def ok_init():
		return True

	server._ensure_initialized = ok_init  # type: ignore[attr-defined]

	result = await server._handle_upload_file({
		"local_path": "C:/does/not/exist.txt",
		"remote_path": "/tmp/sample.txt"
	})

	assert result["status"] == "error"
	assert "Local file not found" in result["message"]


@pytest.mark.asyncio
async def test_download_success(tmp_path):
	server = PentestMCPServer()
	sftp = DummySFTP()
	server.ssh_manager = DummySSHManager(sftp)
	server.tmux_manager = object()

	async def ok_init():
		return True

	server._ensure_initialized = ok_init  # type: ignore[attr-defined]

	local_dest = tmp_path / "nested" / "file.txt"

	result = await server._handle_download_file({
		"remote_path": "/tmp/remote.txt",
		"local_path": str(local_dest)
	})

	assert result["status"] == "success"
	assert local_dest.exists()
	assert ("get", "/tmp/remote.txt", str(local_dest)) in sftp.operations
	assert sftp.closed is True


@pytest.mark.asyncio
async def test_parse_handler_sftp_exit(tmp_path):
	server = PentestMCPServer()
	sftp = DummySFTP()
	server.ssh_manager = DummySSHManager(sftp)
	server.tmux_manager = object()

	async def ok_init():
		return True

	server._ensure_initialized = ok_init  # type: ignore[attr-defined]

	async def fake_parse(tool, fmt, path):
		return {"ok": True, "size": os.path.getsize(path)}

	server._parse_output_file = fake_parse  # type: ignore[attr-defined]

	result = await server._handle_parse_tool_output({
		"tool": "custom",
		"file_path": "/tmp/remote.any",
		"format": "txt"
	})

	assert result["status"] == "success"
	assert sftp.closed is True
