from __future__ import annotations

from types import SimpleNamespace

import pytest

from worai.core import prune_entities_outside_dataset as mod


@pytest.mark.asyncio
async def test_get_dataset_uri_async_success_and_missing(monkeypatch: pytest.MonkeyPatch) -> None:
    class _Client:
        async def __aenter__(self):
            return self

        async def __aexit__(self, *_args):
            return None

    monkeypatch.setattr(mod, "_build_client", lambda *_a, **_k: _Client())

    async def _get_me_ok():
        return SimpleNamespace(dataset_uri="https://d/")

    monkeypatch.setattr(
        mod.wordlift_client,
        "AccountApi",
        lambda _client: SimpleNamespace(get_me=_get_me_ok),
    )
    assert await mod._get_dataset_uri_async("k", "https://api") == "https://d/"

    async def _get_me_missing():
        return SimpleNamespace(dataset_uri=None)

    monkeypatch.setattr(
        mod.wordlift_client,
        "AccountApi",
        lambda _client: SimpleNamespace(get_me=_get_me_missing),
    )
    with pytest.raises(RuntimeError, match="Failed to resolve dataset_uri"):
        await mod._get_dataset_uri_async("k", "https://api")


def test_normalize_and_fetch_outside_dataset(monkeypatch: pytest.MonkeyPatch) -> None:
    assert mod._normalize_dataset_prefixes("https://dataset/") == ["https://dataset", "https://dataset/"]

    seen: dict[str, str] = {}

    def _graphql(endpoint: str, api_key: str, query: str):
        seen["endpoint"] = endpoint
        seen["api_key"] = api_key
        seen["query"] = query
        return {
            "data": {
                "entities": [
                    {"iri": "https://dataset/x"},
                    {"iri": "https://other/y"},
                    {},
                    "bad",
                ]
            }
        }

    monkeypatch.setattr(mod, "graphql_request", _graphql)
    outside = mod.fetch_outside_dataset_iris(
        "https://gql",
        "wl",
        "https://dataset/",
        "query { entities(dataset: \"{dataset_uri}\") { iri } }",
    )
    assert outside == ["https://other/y"]
    assert "https://dataset/" in seen["query"]


def test_fetch_outside_dataset_bad_shape(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(mod, "graphql_request", lambda *_a, **_k: {"data": {}})
    with pytest.raises(RuntimeError, match="Unexpected GraphQL response structure"):
        mod.fetch_outside_dataset_iris("e", "k", "https://d", "q")


def test_run_with_and_without_dataset_lookup(monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]) -> None:
    monkeypatch.setattr(mod, "fetch_outside_dataset_iris", lambda *_a, **_k: ["https://a", "https://b"])
    mod.run(mod.ListOptions(api_key="wl", dataset_uri="https://dataset/", limit=1))
    out = capsys.readouterr().out
    assert "Dataset URI: https://dataset/" in out
    assert "Found 1 entities" in out

    def _fake_run(coro):
        try:
            coro.close()
        except Exception:
            pass
        return "https://from-lookup/"

    monkeypatch.setattr(mod.asyncio, "run", _fake_run)
    mod.run(mod.ListOptions(api_key="wl", dataset_uri=None, limit=0))
    out = capsys.readouterr().out
    assert "Dataset URI: https://from-lookup/" in out
