"""Layer 2 — entity extractor (the local NER brain).

Pulls the people / companies / products / amounts out of a durable fact using
the LOCAL model (Ollama) — the same model the passive extractor already runs,
so it costs nothing extra and never touches an API. The memory service turns
whatever we return into shared :Entity nodes and (:Memory)-[:MENTIONS]->(:Entity)
edges, so "pull up Al Temimi" returns everything hanging off that one node.

Called from memory_control.fire_save() in the background save thread, and by the
one-off entity backfill over memories saved before Layer 2. MEM_ENTITIES=0 off.
"""
from __future__ import annotations

import json
import os
import re
import urllib.request

OLLAMA = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434")
MODEL = os.environ.get("MEM_EXTRACT_MODEL",
                       os.environ.get("LOCAL_MODEL", "qwen3:4b-instruct"))

_PROMPT = """Extract the named entities from this fact about Ahmed's work and \
life. List the specific PEOPLE, COMPANIES/clients, PRODUCTS, PLACES, and \
AMOUNTS (money or notable numbers) it names by their real name. Skip generic \
words like "client", "meeting", "the company", "a lead" unless a real name is \
given.

Output ONLY a JSON array; each item {{"name": "...", "type": "..."}} where type \
is one of: person, company, product, place, amount. If there are no real named \
entities, output exactly [].

Fact: "{text}"

JSON:"""

_TYPES = {"person", "company", "product", "place", "amount"}

# The user himself is in every fact — an "Ahmed" node would be a useless hub
# connecting everything. Skip him and obviously-generic non-names.
_SKIP = {"ahmed", "sir", "jarvis", "me", "him", "he",
         "client", "clients", "a client", "the client", "two clients",
         "customer", "customers", "lead", "a lead", "meeting", "company",
         "the company", "a company", "his company", "employee", "employees"}


def _is_generic(name: str) -> bool:
    n = name.lower().strip()
    if n in _SKIP:
        return True
    # a bare count of clients/leads/etc. ("two clients", "3 leads") isn't a name
    return bool(re.fullmatch(
        r"(a|an|the|one|two|three|four|five|\d+)?\s*"
        r"(client|clients|lead|leads|customer|customers|company|companies|"
        r"meeting|meetings|employee|employees)", n))


def _post(text: str) -> str:
    body = json.dumps({
        "model": MODEL, "stream": False,
        "messages": [{"role": "user", "content": _PROMPT.format(text=text)}],
        "options": {"temperature": 0},
    }).encode()
    req = urllib.request.Request(f"{OLLAMA}/api/chat", data=body,
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=30) as r:
        o = json.loads(r.read())
    return ((o.get("message") or {}).get("content") or "").strip()


def extract(text: str) -> list[dict]:
    """Return [{"name","type"}] entities from `text`; [] on anything unexpected."""
    text = (text or "").strip()
    if os.environ.get("MEM_ENTITIES", "1") == "0" or len(text) < 8:
        return []
    try:
        raw = _post(text)
    except Exception:  # noqa: BLE001 — never disrupt a save
        return []
    m = re.search(r"\[.*\]", raw, re.DOTALL)   # dig the JSON array out
    if not m:
        return []
    try:
        data = json.loads(m.group(0))
    except Exception:  # noqa: BLE001
        return []
    out, seen = [], set()
    for e in data if isinstance(data, list) else []:
        if not isinstance(e, dict):
            continue
        name = str(e.get("name", "")).strip()
        if len(name) < 2 or _is_generic(name):
            continue
        etype = str(e.get("type", "thing")).strip().lower()
        if etype not in _TYPES:
            etype = "thing"
        key = name.lower()
        if key in seen:
            continue
        seen.add(key)
        out.append({"name": name, "type": etype})
    return out[:12]