from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, List, Optional

from pydantic import BaseModel

from sync_state_machine.common.sync_node import SyncNode
from sync_state_machine.datasource.datasource import BaseDataSource
from sync_state_machine.datasource.task_result import TaskResult


class MockSchema(BaseModel):
    id: Optional[str] = None
    name: str = "mock"


class MockSyncNode(SyncNode[MockSchema]):
    node_type = "mock"
    schema = MockSchema


@dataclass
class HandlerScript:
    create: str = "success"  # success | failed | skipped | in_progress
    update: str = "success"
    delete: str = "success"
    poll_sequence: List[str] = field(default_factory=lambda: ["success"])  # success | failed | in_progress


class ScriptedMockHandler:
    def __init__(self, script: HandlerScript):
        self._script = script
        self._collection = None
        self._poll_index = 0

    @property
    def node_type(self) -> str:
        return "mock"

    @property
    def node_class(self):
        return MockSyncNode

    @property
    def schema(self):
        return MockSchema

    def set_collection(self, collection):
        self._collection = collection

    async def load(self):
        return []

    async def create_all(self, nodes):
        return [self._build_result(self._script.create, node.node_id, "create") for node in nodes]

    async def update_all(self, nodes):
        return [
            self._build_result(
                self._script.update,
                node.node_id,
                "update",
            )
            for node in nodes
        ]

    async def delete_all(self, nodes):
        return [self._build_result(self._script.delete, node.node_id, "delete") for node in nodes]

    async def poll_tasks(self, task_ids: List[str]) -> Dict[str, TaskResult]:
        status = self._script.poll_sequence[min(self._poll_index, len(self._script.poll_sequence) - 1)]
        self._poll_index += 1
        out: Dict[str, TaskResult] = {}
        for task_id in task_ids:
            out[task_id] = self._build_poll_result(status, task_id)
        return out

    def _build_result(
        self,
        mode: str,
        node_id: str,
        phase: str,
    ) -> TaskResult:
        if mode == "success":
            return TaskResult.success(node_id=node_id, data_id=f"{phase}-{node_id}")
        if mode == "failed":
            return TaskResult.failed(node_id=node_id, error=f"{phase} failed")
        if mode == "skipped":
            return TaskResult.skipped(
                node_id=node_id,
                reason=f"{phase} skipped",
            )
        if mode == "in_progress":
            return TaskResult.in_progress(node_id=node_id, task_id=f"task-{phase}-{node_id}")
        raise ValueError(f"Unsupported mode: {mode}")

    def _build_poll_result(self, mode: str, task_id: str) -> TaskResult:
        if mode == "success":
            return TaskResult.success(node_id=self._node_id_from_task(task_id), data_id=f"done-{task_id}")
        if mode == "failed":
            return TaskResult.failed(node_id=self._node_id_from_task(task_id), error="poll failed")
        if mode == "in_progress":
            return TaskResult.in_progress(node_id=self._node_id_from_task(task_id), task_id=task_id)
        raise ValueError(f"Unsupported poll mode: {mode}")

    @staticmethod
    def _node_id_from_task(task_id: str) -> str:
        return task_id.split("-")[-1]


class MockDataSource(BaseDataSource):
    pass