"""Diffbot NLP API: entity identification, resolution, and sentiment."""

from typing import TYPE_CHECKING, Any, Dict

if TYPE_CHECKING:
    from .client import Diffbot, DiffbotAsync

NLP_BASE = "https://nl.diffbot.com/v1/"
NLP_FIELDS = "entities,sentiment"


def entities(
    client: "Diffbot",
    text: str,
    *,
    lang: str = "auto",
) -> Dict[str, Any]:
    params = {"token": client.token, "fields": NLP_FIELDS}
    payload = [{"lang": lang, "format": "plain text", "content": text}]
    response = client._http.post(client.nlp_url, params=params, json=payload)
    client._raise_for_status(response)
    data = response.json()
    return data[0] if isinstance(data, list) else data


async def entities_async(
    client: "DiffbotAsync",
    text: str,
    *,
    lang: str = "auto",
) -> Dict[str, Any]:
    params = {"token": client.token, "fields": NLP_FIELDS}
    payload = [{"lang": lang, "format": "plain text", "content": text}]
    response = await client._http.post(client.nlp_url, params=params, json=payload)
    client._raise_for_status(response)
    data = response.json()
    return data[0] if isinstance(data, list) else data