from __future__ import annotations

from pathlib import Path

import pytest
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.config import (
    BaseDataSourceConfig,
    JsonlDataSourceConfig,
    LoggingConfig,
    PersistConfig,
    PipelineRunConfig,
    register_datasource_type,
)
from sync_state_machine.datasource import BaseDataSource
from sync_state_machine.datasource.handler import BaseNodeHandler
from sync_state_machine.datasource.jsonl import BaseJsonlHandler, JsonlDataSource
from sync_state_machine.pipeline.factory import create_pipeline_from_config
from sync_state_machine.sync_system.strategy import DefaultSyncStrategy


class TestRegisteredDbConfig(BaseDataSourceConfig):
    __test__ = False
    type: str = "test_db"


class TestDbSchema(BaseModel):
    __test__ = False
    id: str
    name: str


class TestDbNode(SyncNode[TestDbSchema]):
    __test__ = False
    node_type = "test_db_handler"
    schema = TestDbSchema


class TestDbStrategy(DefaultSyncStrategy[TestDbSchema]):
    __test__ = False
    schema = TestDbSchema


class TestDbHandler(BaseNodeHandler[TestDbSchema]):
    __test__ = False

    def __init__(self, datasource: BaseDataSource):
        super().__init__("test_db_handler", TestDbNode, TestDbSchema)
        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 TestJsonlHandler(BaseJsonlHandler):
    __test__ = False

    def __init__(self, datasource: JsonlDataSource):
        super().__init__(datasource, "test_db_handler", TestDbNode, TestDbSchema)


def _register_test_db_datasource_type() -> None:
    register_datasource_type(
        "test_db",
        config_model=TestRegisteredDbConfig,
        factory=lambda config, endpoint_name: BaseDataSource(),
        replace=True,
    )


@pytest.mark.asyncio
async def test_register_handlers_uses_db_handler_for_db_datasource(tmp_path: Path) -> None:
    _register_test_db_datasource_type()

    if not DomainRegistry.is_registered("test_db_handler"):
        DomainRegistry.register(
            node_type="test_db_handler",
            schema=TestDbSchema,
            node_class=TestDbNode,
            strategy_class=TestDbStrategy,
            jsonl_handler_class=TestJsonlHandler,
        )
    DomainRegistry.register_handler("test_db_handler", "test_db", TestDbHandler)

    config = PipelineRunConfig(
        node_types=["test_db_handler"],
        local_datasource=TestRegisteredDbConfig(type="test_db"),
        remote_datasource=JsonlDataSourceConfig(
            type="jsonl",
            jsonl_dir=str(tmp_path),
        ),
        strategies=[],
        persist=PersistConfig(db_path=str(tmp_path / "pipeline.db")),
        logging=LoggingConfig(initialize=False),
    )

    pipeline = await create_pipeline_from_config(config)
    try:
        assert isinstance(pipeline.local_datasource.get_handler("test_db_handler"), TestDbHandler)
        assert isinstance(pipeline.remote_datasource.get_handler("test_db_handler"), TestJsonlHandler)
    finally:
        await pipeline.close()