from __future__ import annotations

import json

import pytest
from pydantic import BaseModel

from sync_state_machine.common.collection import DataCollection
from sync_state_machine.common.sync_node import SyncNode
from sync_state_machine.datasource.jsonl.datasource import JsonlDataSource
from sync_state_machine.datasource.jsonl.handler import BaseJsonlHandler


class _ProjectSchema(BaseModel):
    id: str


class _BizSchema(BaseModel):
    id: str
    project_id: str | None = None


class _ProjectNode(SyncNode[_ProjectSchema]):
    node_type = "project"
    schema = _ProjectSchema


class _BizNode(SyncNode[_BizSchema]):
    node_type = "contract"
    schema = _BizSchema


@pytest.mark.asyncio
async def test_jsonl_project_load_applies_target_project_filter(tmp_path):
    data_file = tmp_path / "project_1.jsonl"
    data_file.write_text(
        "\n".join([
            json.dumps({"id": "p1"}, ensure_ascii=False),
            json.dumps({"id": "p2"}, ensure_ascii=False),
        ])
        + "\n",
        encoding="utf-8",
    )

    datasource = JsonlDataSource(tmp_path)
    handler = BaseJsonlHandler(datasource, "project", _ProjectNode, _ProjectSchema)
    handler.set_handler_config({"data_id_filter": ["p1"]})
    datasource.add_handler(handler)
    datasource.set_collection(DataCollection("local"))

    await datasource.load_all(order=["project"])

    loaded = datasource._collection.filter(node_type="project")  # type: ignore[union-attr]
    assert {node.data_id for node in loaded} == {"p1"}


@pytest.mark.asyncio
async def test_jsonl_non_project_load_filters_rows_by_project_id(tmp_path):
    data_file = tmp_path / "contract_1.jsonl"
    data_file.write_text(
        "\n".join([
            json.dumps({"id": "c1", "project_id": "p1"}, ensure_ascii=False),
            json.dumps({"id": "c2", "project_id": "p2"}, ensure_ascii=False),
            json.dumps({"id": "c3"}, ensure_ascii=False),
        ])
        + "\n",
        encoding="utf-8",
    )

    datasource = JsonlDataSource(tmp_path)
    handler = BaseJsonlHandler(datasource, "contract", _BizNode, _BizSchema)
    handler.set_handler_config({"data_id_filter": ["p1"]})
    datasource.add_handler(handler)
    datasource.set_collection(DataCollection("local"))

    await datasource.load_all(order=["contract"])

    loaded = datasource._collection.filter(node_type="contract")  # type: ignore[union-attr]
    assert {node.data_id for node in loaded} == {"c1", "c3"}


@pytest.mark.asyncio
async def test_jsonl_non_project_load_inherits_project_filter_from_loaded_projects(tmp_path):
    project_file = tmp_path / "project_1.jsonl"
    project_file.write_text(
        "\n".join([
            json.dumps({"id": "p1"}, ensure_ascii=False),
        ])
        + "\n",
        encoding="utf-8",
    )
    contract_file = tmp_path / "contract_1.jsonl"
    contract_file.write_text(
        "\n".join([
            json.dumps({"id": "c1", "project_id": "p1"}, ensure_ascii=False),
            json.dumps({"id": "c2", "project_id": "p2"}, ensure_ascii=False),
            json.dumps({"id": "c3"}, ensure_ascii=False),
        ])
        + "\n",
        encoding="utf-8",
    )

    datasource = JsonlDataSource(tmp_path)
    project_handler = BaseJsonlHandler(datasource, "project", _ProjectNode, _ProjectSchema)
    project_handler.set_handler_config({"data_id_filter": ["p1"]})
    contract_handler = BaseJsonlHandler(datasource, "contract", _BizNode, _BizSchema)
    datasource.add_handler(project_handler)
    datasource.add_handler(contract_handler)
    datasource.set_collection(DataCollection("local"))

    await datasource.load_all(order=["project", "contract"])

    loaded = datasource._collection.filter(node_type="contract")  # type: ignore[union-attr]
    assert {node.data_id for node in loaded} == {"c1", "c3"}
