from pathlib import Path

from agentci.domain.models import Job, JobKind, Workflow, WorkflowKind
from agentci.worker import Worker


class FakeStorage:
    def __init__(self, workflow: Workflow | None) -> None:
        self.workflow = workflow

    async def get_workflow(self, _workflow_id: str) -> Workflow | None:
        return self.workflow


class FakeOpenCode:
    def __init__(self) -> None:
        self.aborted: set[tuple[str, Path]] = set()

    async def abort(self, session_id: str, workspace: Path) -> None:
        self.aborted.add((session_id, workspace))


def job(*, workflow_id: str | None, runtime_session_id: str | None = None) -> Job:
    return Job(
        id="job",
        kind=JobKind.FIX,
        target_key="org/repo:pr:1",
        repo_owner="org",
        repo_name="repo",
        issue_number=1,
        pr_number=1,
        requester="alice",
        message="",
        comment_id=1,
        workflow_id=workflow_id,
        runtime_session_id=runtime_session_id,
    )


def worker(tmp_path: Path, storage: FakeStorage, opencode: FakeOpenCode) -> Worker:
    return Worker(
        storage=storage,  # type: ignore[arg-type]
        gitea=None,  # type: ignore[arg-type]
        opencode=opencode,  # type: ignore[arg-type]
        dispatcher=None,  # type: ignore[arg-type]
        poll_seconds=1,
        workspaces_dir=tmp_path,
    )


async def test_recovery_aborts_all_workflow_sessions(tmp_path: Path) -> None:
    workspace = tmp_path / "workflow" / "repo"
    workflow = Workflow(
        id="flow",
        kind=WorkflowKind.IMPLEMENT,
        repo_owner="org",
        repo_name="repo",
        issue_number=1,
        workspace_path=workspace,
        base_sha="base",
        primary_session_id="primary",
        reviewer_session_id="reviewer",
    )
    opencode = FakeOpenCode()

    await worker(tmp_path, FakeStorage(workflow), opencode)._abort_job_sessions(
        job(workflow_id=workflow.id)
    )

    assert opencode.aborted == {("primary", workspace), ("reviewer", workspace)}


async def test_recovery_aborts_one_shot_fix_session(tmp_path: Path) -> None:
    opencode = FakeOpenCode()

    await worker(tmp_path, FakeStorage(None), opencode)._abort_job_sessions(
        job(workflow_id=None, runtime_session_id="fix-session")
    )

    assert opencode.aborted == {("fix-session", tmp_path / "fix-job" / "repo")}