from __future__ import annotations

from typing import Any

import pytest
from pydantic import BaseModel

from sync_state_machine.common.sync_node import SyncNode
from sync_state_machine.domain.project.api_handler import ProjectApiHandler
from sync_state_machine.datasource.task_result import TaskStatus


class _ProjectUpdateSchema(BaseModel):
    id: str
    dynamic_investment: float


class _ProjectUpdateNode(SyncNode[_ProjectUpdateSchema]):
    node_type = "project"
    schema = _ProjectUpdateSchema


class _MockApiClient:
    def __init__(self) -> None:
        self.requests: list[dict[str, Any]] = []
        self._poll_count = 0

    async def request(self, method: str, endpoint: str, **kwargs: Any) -> dict[str, Any]:
        self.requests.append({"method": method, "endpoint": endpoint, **kwargs})
        return {"code": 200, "message": "OK", "data": {"biz_id": kwargs.get("data", {}).get("biz_id")}}

    async def get_push_logs(self, task_ids: list[str]) -> dict[str, dict[str, Any]]:
        self._poll_count += 1
        if self._poll_count == 1:
            return {task_ids[0]: {}}
        return {task_ids[0]: {"code": 200, "data": {"biz_id": "project-1"}}}


@pytest.mark.asyncio
async def test_project_update_waits_for_async_push_completion() -> None:
    handler = ProjectApiHandler()
    handler.set_handler_config({"data_id_filter": ["project-1"]})
    handler.api_client = _MockApiClient()

    node = _ProjectUpdateNode(
        node_id="project-node-1",
        data_id="project-1",
        data={"id": "project-1", "dynamic_investment": 200.0},
        origin_data={"id": "project-1", "dynamic_investment": 100.0},
    )

    results = await handler.update_all([node])

    assert len(results) == 1
    assert results[0].status == TaskStatus.SUCCESS
    assert handler.api_client.requests[0]["endpoint"] == "/project/dynamic_investment"
    assert handler.api_client._poll_count == 2
