from __future__ import annotations

from collections import defaultdict
from typing import Dict, List, Optional, Sequence

from pydantic import BaseModel

from sync_state_machine.common.registry import DomainRegistry
from sync_state_machine.common.sync_node import SyncNode
from sync_state_machine.datasource.datasource import BaseDataSource
from sync_state_machine.datasource.handler import ValidationResult
from sync_state_machine.datasource.task_result import TaskResult
from sync_state_machine.sync_system.strategy import DefaultSyncStrategy


class TestProjectSchema(BaseModel):
    __test__ = False
    id: str
    name: str
    code: str


class TestContractSchema(BaseModel):
    __test__ = False
    id: str
    name: str
    code: str
    project_ref: str


class TestProjectNode(SyncNode[TestProjectSchema]):
    __test__ = False
    node_type = "test_project"
    schema = TestProjectSchema


class TestContractNode(SyncNode[TestContractSchema]):
    __test__ = False
    node_type = "test_contract"
    schema = TestContractSchema


class TestProjectStrategy(DefaultSyncStrategy[TestProjectSchema]):
    __test__ = False
    schema = TestProjectSchema


class TestContractStrategy(DefaultSyncStrategy[TestContractSchema]):
    __test__ = False
    schema = TestContractSchema


def register_test_domain() -> None:
    if not DomainRegistry.is_registered("test_project"):
        DomainRegistry.register(
            node_type="test_project",
            schema=TestProjectSchema,
            node_class=TestProjectNode,
            strategy_class=TestProjectStrategy,
        )
    if not DomainRegistry.is_registered("test_contract"):
        DomainRegistry.register(
            node_type="test_contract",
            schema=TestContractSchema,
            node_class=TestContractNode,
            strategy_class=TestContractStrategy,
        )


def build_project_node(
    node_id: str,
    data_id: str,
    *,
    name: str,
    code: str,
    data_present: bool = True,
):
    data = {"id": data_id, "name": name, "code": code} if data_present else None
    return TestProjectNode(node_id=node_id, data_id=data_id, data=data)


def build_contract_node(
    node_id: str,
    data_id: str,
    *,
    name: str,
    code: str,
    project_ref: str,
    data_present: bool = True,
):
    data = (
        {
            "id": data_id,
            "name": name,
            "code": code,
            "project_ref": project_ref,
        }
        if data_present
        else None
    )
    return TestContractNode(node_id=node_id, data_id=data_id, data=data)


class InMemoryNodeHandler:
    def __init__(
        self,
        *,
        node_type: str,
        node_class,
        schema,
        load_nodes: Sequence[SyncNode],
        create_mode_by_code: Optional[Dict[str, str]] = None,
        update_mode_by_code: Optional[Dict[str, str]] = None,
    ):
        self._node_type = node_type
        self._node_class = node_class
        self._schema = schema
        self._load_nodes = load_nodes
        self._collection = None
        self._create_mode_by_code = create_mode_by_code or {}
        self._update_mode_by_code = update_mode_by_code or {}
        self._poll_sequences: Dict[str, List[str]] = defaultdict(list)

    @property
    def node_type(self) -> str:
        return self._node_type

    @property
    def node_class(self):
        return self._node_class

    @property
    def schema(self):
        return self._schema

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

    def extract_id(self, data: Dict[str, str]) -> str:
        return str(data.get("id", ""))

    async def validate(self, data: Dict[str, str]) -> ValidationResult:
        try:
            self._schema.model_validate(data)
            return ValidationResult(is_valid=True)
        except Exception as exc:
            return ValidationResult(is_valid=False, errors=[str(exc)])

    def _create_node(self, data: Dict[str, str]):
        data_id = self.extract_id(data)
        return self._node_class(
            node_id=f"{self._node_type}-{data_id}",
            data_id=data_id,
            data=data,
        )

    async def load(self):
        return [
            self._node_class(
                node_id=node.node_id,
                data_id=node.data_id,
                data=node.get_data(),
                origin_data=node.get_data(),
                action=node.action,
                status=node.status,
                binding_status=node.binding_status,
                depend_ids=list(node.depend_ids),
                context=dict(node.context),
            )
            for node in self._load_nodes
        ]

    async def create_all(self, nodes):
        results = []
        for node in nodes:
            data = node.get_data() or {}
            mode = self._create_mode_by_code.get(data.get("code", ""), "success")
            results.append(self._result_for_mode(mode, node.node_id, node_code=data.get("code", ""), phase="create"))
        return results

    async def update_all(self, nodes):
        results = []
        for node in nodes:
            data = node.get_data() or {}
            mode = self._update_mode_by_code.get(data.get("code", ""), "success")
            results.append(self._result_for_mode(mode, node.node_id, node_code=data.get("code", ""), phase="update"))
        return results

    async def delete_all(self, nodes):
        return [TaskResult.success(node_id=node.node_id, data_id=node.data_id) for node in nodes]

    async def poll_tasks(self, task_ids: List[str]) -> Dict[str, TaskResult]:
        out: Dict[str, TaskResult] = {}
        for task_id in task_ids:
            sequence = self._poll_sequences.get(task_id) or ["success"]
            mode = sequence.pop(0)
            if not sequence:
                sequence = [mode]
            self._poll_sequences[task_id] = sequence
            node_id = task_id.split("::", 1)[1]
            if mode == "success":
                out[task_id] = TaskResult.success(node_id=node_id, data_id=f"polled-{node_id}")
            elif mode == "failed":
                out[task_id] = TaskResult.failed(node_id=node_id, error="poll failed")
            else:
                out[task_id] = TaskResult.in_progress(node_id=node_id, task_id=task_id)
        return out

    def _result_for_mode(self, mode: str, node_id: str, *, node_code: str, phase: str) -> TaskResult:
        if mode == "success":
            return TaskResult.success(node_id=node_id, data_id=f"{phase}-ok-{node_id}")
        if mode == "failed":
            return TaskResult.failed(node_id=node_id, error=f"{phase} failed for {node_code}")
        if mode == "skipped":
            return TaskResult.skipped(node_id=node_id, reason=f"{phase} skipped for {node_code}")
        if mode == "in_progress_success":
            task_id = f"task::{node_id}"
            self._poll_sequences[task_id] = ["success"]
            return TaskResult.in_progress(node_id=node_id, task_id=task_id)
        if mode == "in_progress_failed":
            task_id = f"task::{node_id}"
            self._poll_sequences[task_id] = ["failed"]
            return TaskResult.in_progress(node_id=node_id, task_id=task_id)
        raise ValueError(f"Unsupported mode: {mode}")


class InMemoryDataSource(BaseDataSource):
    pass