from __future__ import annotations

from pathlib import Path

import pytest
from pydantic import BaseModel, ConfigDict

from sync_state_machine.common.registry import DomainRegistry
from sync_state_machine.common.sync_node import SyncNode
from sync_state_machine.config import (
    BaseDataSourceConfig,
    LoggingConfig,
    PersistConfig,
    PipelineRunConfig,
    build_config,
    get_datasource_factory,
    register_datasource_type,
)
from sync_state_machine.datasource import BaseDataSource
from sync_state_machine.datasource import ensure_builtin_datasource_types_registered
from sync_state_machine.datasource.handler import BaseNodeHandler
from sync_state_machine.pipeline.factory import create_pipeline_from_config
from sync_state_machine.sync_system.strategy import DefaultSyncStrategy


class MemoryConfig(BaseDataSourceConfig):
    model_config = ConfigDict(extra="forbid", validate_assignment=True)

    type: str = "memory"
    namespace: str = "default"


class MemorySchema(BaseModel):
    __test__ = False

    id: str
    name: str | None = None


class MemoryNode(SyncNode[MemorySchema]):
    __test__ = False

    node_type = "memory_demo"
    schema = MemorySchema


class MemoryStrategy(DefaultSyncStrategy[MemorySchema]):
    __test__ = False

    schema = MemorySchema


class MemoryDataSource(BaseDataSource):
    def __init__(self, namespace: str):
        super().__init__()
        self.namespace = namespace


class MemoryHandler(BaseNodeHandler[MemorySchema]):
    __test__ = False

    def __init__(self, datasource: MemoryDataSource):
        super().__init__("memory_demo", MemoryNode, MemorySchema)
        self.datasource = datasource

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

    async def load(self):
        return []

    async def create_all(self, nodes):
        return []

    async def update_all(self, nodes):
        return []


class MemoryRemoteHandler(BaseNodeHandler[MemorySchema]):
    __test__ = False

    def __init__(self, datasource: BaseDataSource):
        super().__init__("memory_demo", MemoryNode, MemorySchema)
        self.datasource = datasource

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

    async def load(self):
        return []

    async def create_all(self, nodes):
        return []

    async def update_all(self, nodes):
        return []


def _ensure_memory_domain_registered() -> None:
    if DomainRegistry.is_registered("memory_demo"):
        reg = DomainRegistry.get_registration("memory_demo")
        reg.schema = MemorySchema
        reg.node_class = MemoryNode
        reg.strategy_class = MemoryStrategy
        reg.handler_classes["memory_remote"] = MemoryRemoteHandler
        reg.handler_classes["memory"] = MemoryHandler
        return

    DomainRegistry.register(
        node_type="memory_demo",
        schema=MemorySchema,
        node_class=MemoryNode,
        strategy_class=MemoryStrategy,
    )
    DomainRegistry.register_handler("memory_demo", "memory", MemoryHandler)
    DomainRegistry.register_handler("memory_demo", "memory_remote", MemoryRemoteHandler)


def _memory_factory(config: MemoryConfig, endpoint_name: str) -> MemoryDataSource:
    return MemoryDataSource(namespace=f"{endpoint_name}:{config.namespace}")


def test_builtin_datasource_factories_are_registered() -> None:
    ensure_builtin_datasource_types_registered(replace=True)

    assert get_datasource_factory("jsonl") is not None
    assert get_datasource_factory("api") is not None


@pytest.mark.asyncio
async def test_create_pipeline_from_config_supports_registered_custom_datasource_type(tmp_path: Path) -> None:
    register_datasource_type("memory", config_model=MemoryConfig, factory=_memory_factory, replace=True)
    _ensure_memory_domain_registered()

    config = PipelineRunConfig(
        node_types=["memory_demo"],
        local_datasource=MemoryConfig(type="memory", namespace="alpha"),
        remote_datasource=MemoryConfig(type="memory", namespace="omega"),
        strategies=[],
        persist=PersistConfig(db_path=str(tmp_path / "memory-registry.db")),
        logging=LoggingConfig(initialize=False),
    )

    pipeline = await create_pipeline_from_config(config)
    try:
        assert isinstance(pipeline.local_datasource, MemoryDataSource)
        assert pipeline.local_datasource.namespace == "local:alpha"
        assert isinstance(pipeline.remote_datasource, MemoryDataSource)
        assert pipeline.remote_datasource.namespace == "remote:omega"
        assert isinstance(pipeline.local_datasource.get_handler("memory_demo"), MemoryHandler)
        assert isinstance(pipeline.remote_datasource.get_handler("memory_demo"), MemoryHandler)
    finally:
        await pipeline.close()


def test_build_config_supports_custom_datasource_type_loading(tmp_path: Path) -> None:
    register_datasource_type("memory_bootstrap", config_model=MemoryConfig, factory=_memory_factory, replace=True)

    profile_path = tmp_path / "custom-memory.yaml"
    profile_path.write_text(
        "node_types: []\n"
        "local_datasource:\n"
        "  type: memory_bootstrap\n"
        "  namespace: local-demo\n"
        "remote_datasource:\n"
        "  type: jsonl\n"
        "  jsonl_dir: ${PROJECT_ROOT}/remote_datasource/datasource\n"
        "persist:\n"
        "  db_path: ${PROJECT_ROOT}/runtime/test-memory-bootstrap.db\n"
        "logging:\n"
        "  initialize: false\n",
        encoding="utf-8",
    )

    config = build_config(project_root=tmp_path, overrides_file=profile_path)

    assert config.local_datasource.type == "memory_bootstrap"
    assert getattr(config.local_datasource, "namespace") == "local-demo"
    assert config.remote_datasource.type == "jsonl"
