"""CLI command: db entities — identify and resolve entities in text via the Diffbot NLP API."""

import json

import sys

import click

from rich.console import Console

from rich.markup import escape

from rich.progress import Progress, SpinnerColumn, TextColumn

from rich.table import Table

from diffbot import AuthError, APIError

from ._common import get_client

console = Console()

_PALETTE = [

"cyan", "magenta", "yellow", "green", "blue", "red",

"bright_cyan", "bright_magenta", "bright_yellow", "bright_green",

]

def _sentiment_label(score) -> str:

if score is None:

return ""

if score > 0.1:

return f"[green]+{score:.2f}[/green]"

if score < -0.1:

return f"[red]{score:.2f}[/red]"

return f"[dim]{score:.2f}[/dim]"

def _highlight_mentions(text: str, ents: list, palette: list) -> str:

spans = []

for i, e in enumerate(ents):

color = palette[i % len(palette)]

for m in e.get("mentions", []):

begin = m.get("beginOffset")

end = m.get("endOffset")

if begin is not None and end is not None:

spans.append((begin, end, color))

if not spans:

return escape(text)

spans.sort(key=lambda s: s[0])

parts = []

pos = 0

for begin, end, color in spans:

if begin < pos:

continue # skip overlapping spans

parts.append(escape(text[pos:begin]))

parts.append(f"[bold underline {color}]{escape(text[begin:end])}[/bold underline {color}]")

pos = end

parts.append(escape(text[pos:]))

return "".join(parts)

@click.command(name="entities")

@click.argument("text", required=False)

@click.option("-f", "--format", "fmt", type=click.Choice(["table", "json", "dql"]), default="table",

help="Output format: table (default), json (raw), dql (id:or(...) filter ready for db dql export)")

@click.option("--lang", default="auto", help="Language code (default: auto)")

def entities(text, fmt, lang):

"""Identify and resolve entities and sentiment in text using the Diffbot NLP API.

TEXT can be passed as an argument or piped via stdin.

Entity IDs returned can be used in DQL for fast lookups:

db entities "..." -f dql | xargs db dql export

"""

if text is None:

if sys.stdin.isatty():

raise click.UsageError("provide TEXT as an argument or pipe text via stdin")

text = sys.stdin.read().strip()

if not text:

raise click.UsageError("no text provided")

db = get_client()

try:

is_interactive = sys.stdout.isatty()

if is_interactive:

with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True) as progress:

progress.add_task(description="Resolving entities...", total=None)

result = db.entities(text, lang=lang)

else:

result = db.entities(text, lang=lang)

if fmt == "json":

sys.stdout.write(json.dumps(result, indent=2))

sys.stdout.write("\n")

return

ents = result.get("entities", [])

if fmt == "dql":

ids = []

for e in ents:

raw = e.get("id") or e.get("diffbotUri") or ""

eid = raw.rstrip("/").rsplit("/", 1)[-1] if raw else ""

if eid:

ids.append(eid)

if not ids:

return

sys.stdout.write('id:or(' + ",".join(f'"{i}"' for i in ids) + ')\n')

return

doc_sentiment = result.get("sentiment")

if doc_sentiment is not None:

console.print(f"Document sentiment: {_sentiment_label(doc_sentiment)}")

console.print()

if not ents:

console.print("[dim]No entities found.[/dim]")

return

table = Table(show_header=True, header_style="bold", show_lines=False)

table.add_column("Entity", overflow="fold")

table.add_column("Type", no_wrap=True)

table.add_column("Confidence", no_wrap=True, justify="right")

table.add_column("Salience", no_wrap=True, justify="right")

table.add_column("Sentiment", no_wrap=True, justify="right")

table.add_column("Diffbot ID", overflow="fold", style="dim")

for i, e in enumerate(ents):

color = _PALETTE[i % len(_PALETTE)]

name = e.get("name") or e.get("label") or ""

all_types = e.get("allTypes") or []

etype = all_types[0]["name"] if all_types else (e.get("type") or "")

confidence = e.get("confidence")

salience = e.get("salience")

sentiment = e.get("sentiment")

raw = e.get("id") or e.get("diffbotUri") or ""

eid = raw.rstrip("/").rsplit("/", 1)[-1] if raw else ""

conf_str = f"{confidence:.2f}" if confidence is not None else ""

sal_str = f"{salience:.2f}" if salience is not None else ""

table.add_row(

f"[{color}]● {name}[/{color}]",

etype, conf_str, sal_str, _sentiment_label(sentiment), eid,

)

console.print(_highlight_mentions(text, ents, _PALETTE))

console.print()

console.print(table)

console.print(f"[dim]{len(ents)} entit{'y' if len(ents) == 1 else 'ies'} found[/dim]")

except AuthError:

click.echo("Error: Invalid or unauthorized API token.", err=True)

raise click.Abort()

except APIError as e:

click.echo(f"API error {e.status_code}: {e.message or e.body}", err=True)

raise click.Abort()