"""C2Connector factory + payload mapping tests."""

from __future__ import annotations

import pytest

from mimic.connectors import (
    C2Connector,
    ConnectorFactory,
    Payload,
    TaskHandle,
    TaskResult,
    TaskStatus,
    UnsupportedPayloadType,
    register_connector,
)
from mimic.connectors.factory import _REGISTRY
from mimic.connectors.payload_map import resolve_native, supports
from mimic.db.models.host import Host
from mimic.db.types import C2Type, PayloadType


@pytest.fixture(autouse=True)
def _clear_registry() -> None:
    """Each test starts with a clean connector registry."""
    snapshot = dict(_REGISTRY)
    _REGISTRY.clear()
    yield
    _REGISTRY.clear()
    _REGISTRY.update(snapshot)


class _NullConnector(C2Connector):
    def authenticate(self, config: dict[str, object]) -> None: ...
    def list_hosts(self, engagement_id: str) -> list[Host]:
        return []

    def execute_task(self, host: Host, payload: Payload) -> TaskHandle:
        return TaskHandle(
            c2=self.name,
            c2_task_id="t-1",
            host_id="h-1",
            payload_type=payload.payload_type,
        )

    def get_task_result(self, handle: TaskHandle) -> TaskResult:
        return TaskResult(status=TaskStatus.COMPLETED, output_text="ok")

    def cancel_task(self, handle: TaskHandle) -> None: ...

    def execute_cleanup(
        self, host: Host, resolved_command: str, params: dict[str, object]
    ) -> TaskResult:
        return TaskResult(status=TaskStatus.COMPLETED)


def test_register_and_build() -> None:
    register_connector(C2Type.MYTHIC)(_NullConnector)
    factory = ConnectorFactory(config_resolver=lambda _: {})
    connector = factory.build(C2Type.MYTHIC)
    assert isinstance(connector, _NullConnector)
    assert connector.name is C2Type.MYTHIC


def test_double_registration_rejected() -> None:
    register_connector(C2Type.MYTHIC)(_NullConnector)
    with pytest.raises(RuntimeError, match="already registered"):
        register_connector(C2Type.MYTHIC)(_NullConnector)


def test_build_unknown_raises_not_implemented() -> None:
    factory = ConnectorFactory(config_resolver=lambda _: {})
    with pytest.raises(NotImplementedError):
        factory.build(C2Type.HOME)


def test_task_status_is_terminal() -> None:
    assert TaskStatus.COMPLETED.is_terminal
    assert TaskStatus.FAILED.is_terminal
    assert TaskStatus.CANCELED.is_terminal
    assert not TaskStatus.RUNNING.is_terminal


def test_mythic_mapping_covers_powershell() -> None:
    assert resolve_native(C2Type.MYTHIC, PayloadType.POWERSHELL) == "powershell"
    assert supports(C2Type.MYTHIC, PayloadType.SHELLCODE)


def test_home_mapping_empty_until_pr2() -> None:
    assert not supports(C2Type.HOME, PayloadType.CMD)
    with pytest.raises(UnsupportedPayloadType):
        resolve_native(C2Type.HOME, PayloadType.CMD)