extraction/spacy_extractor.py

from __future__ import annotations

import re
import threading
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

from logger import get_logger

logger = get_logger(__name__)

# ---------------------------------------------------------------------------
# Type aliases
# ---------------------------------------------------------------------------

# A "schema node" is one of:
#   • a leaf rule dict  → has a "source_type" key  (str leaf)
#   • an array rule     → has "type": "array"
#   • an object rule    → has "type": "object"
# Results mirror the shape: str | list | dict | None at any depth.

SchemaNode = Dict[str, Any]
ResultNode = Union[str, List[Any], Dict[str, Any], None]

VALID_SPACY_LABELS: Dict[str, str] = {
    "ORG":         "Companies, agencies, institutions",
    "PERSON":      "People, including fictional",
    "DATE":        "Absolute or relative dates or periods",
    "MONEY":       "Monetary values, including unit",
    "GPE":         "Countries, cities, states",
    "LOC":         "Non-GPE locations, mountain ranges, bodies of water",
    "PRODUCT":     "Objects, vehicles, foods, etc.",
    "EVENT":       "Named hurricanes, battles, wars, sports events",
    "CARDINAL":    "Numerals that do not fall under another type",
    "PERCENT":     "Percentage, including '%'",
    "QUANTITY":    "Measurements, as of weight or distance",
    "TIME":        "Times smaller than a day",
    "NORP":        "Nationalities or religious or political groups",
    "FAC":         "Buildings, airports, highways, bridges",
    "WORK_OF_ART": "Titles of books, songs, etc.",
    "LAW":         "Named documents made into laws",
    "LANGUAGE":    "Any named language",
    "ORDINAL":     "'first', 'second', etc.",
}


_WHITESPACE_RE = re.compile(r"\s+")
_CURRENCY_RE = re.compile(r"[$€£¥₹]")
_NON_NUMERIC_RE = re.compile(r"[^\d.]")
_DATE_SEP_RE = re.compile(r"[/.]")

# ---------------------------------------------------------------------------
# spaCy singleton
# ---------------------------------------------------------------------------

_nlp_lock = threading.Lock()
_nlp: Any = None


def _get_nlp() -> Any:
    """Return the shared spaCy pipeline, initialising it on first call."""
    global _nlp
    if _nlp is not None:
        return _nlp
    with _nlp_lock:
        if _nlp is None:
            import spacy
            _nlp = spacy.load(
                "en_core_web_sm",
                exclude=["tagger", "parser", "lemmatizer", "attribute_ruler"],
            )
            logger.info("spaCy en_core_web_sm loaded (singleton)")
    return _nlp


def clean_text(text: str) -> str:
    """Normalise whitespace on raw text before handing it to spaCy.

    Defined as a module-level function so it is not shadowed by local
    variables named `cleaned` in the public extract_* functions.
    """
    return _WHITESPACE_RE.sub(" ", text).strip()


# ---------------------------------------------------------------------------
# Normalizer registry
# ---------------------------------------------------------------------------

_normalizer_lock = threading.Lock()
_NORMALIZERS: Dict[str, Callable[[str], str]] = {
    "strip":               lambda s: s.strip(),
    "upper":               lambda s: s.upper(),
    "lower":               lambda s: s.lower(),
    "remove_commas":       lambda s: s.replace(",", ""),
    "remove_spaces":       lambda s: s.replace(" ", ""),
    "remove_newlines":     lambda s: s.replace("\n", " ").replace("\r", ""),
    "collapse_whitespace": lambda s: _WHITESPACE_RE.sub(" ", s).strip(),
    "remove_currency":     lambda s: _CURRENCY_RE.sub("", s),
    "remove_non_numeric":  lambda s: _NON_NUMERIC_RE.sub("", s),
    "normalize_date_sep":  lambda s: _DATE_SEP_RE.sub("-", s),
}


def register_normalizer(name: str, fn: Callable[[str], str]) -> None:
    """Register a custom normalizer. Thread-safe, overwrites silently."""
    with _normalizer_lock:
        _NORMALIZERS[name] = fn


def _apply_normalizers(value: Optional[str], normalize: Any) -> Optional[str]:
    if not isinstance(value, str):
        return None
    if not normalize:
        return value
    if isinstance(normalize, str):
        normalize = [normalize]
    for key in normalize:
        fn = _NORMALIZERS.get(key)
        if fn is None:
            logger.warning("Unknown normalizer %r — skipped", key)
            continue
        try:
            value = fn(value)
        except Exception as exc:
            logger.error("Normalizer %r raised on value %r: %s", key, value, exc)
    return value if value else None


# ---------------------------------------------------------------------------
# Resolver registry
# ---------------------------------------------------------------------------

_resolver_lock = threading.Lock()
_RESOLVERS: Dict[str, Callable[[Dict[str, Any], Any, str], Optional[str]]] = {}


def register_resolver(
    source_type: str,
    fn: Callable[[Dict[str, Any], Any, str], Optional[str]],
) -> None:
    """Register a custom resolver for a source_type. Thread-safe."""
    with _resolver_lock:
        _RESOLVERS[source_type] = fn


# ---------------------------------------------------------------------------
# Regex resolver
# ---------------------------------------------------------------------------

def _build_flags(rule: Dict[str, Any]) -> int:
    flags = 0
    for name in rule.get("flags", []):
        obj = getattr(re, name.upper(), None)
        if obj is None:
            logger.warning("Unknown re flag %r — skipped", name)
            continue
        flags |= obj
    return flags


def _try_group(match: re.Match, capture_group: Any) -> Tuple[bool, Optional[str]]:
    try:
        return True, match.group(capture_group)
    except (IndexError, re.error):
        logger.warning(
            "Group %r does not exist in pattern %r",
            capture_group, match.re.pattern,
        )
        return False, None


def _resolve_regex(rule: Dict[str, Any], text: str) -> Optional[str]:
    primary = rule.get("pattern", "")
    if not primary:
        logger.warning("Regex rule missing 'pattern': %s", rule)
        return None

    flags         = _build_flags(rule)
    capture_group = rule.get("capture_group", 0)
    match_index   = rule.get("match_index", 0)
    normalize     = rule.get("normalize", "")
    strip_chars   = rule.get("strip_chars", "")
    fallbacks     = rule.get("fallback_patterns", [])

    for pat in (primary, *fallbacks):
        try:
            matches = list(re.finditer(pat, text, flags))
        except re.error as exc:
            logger.error("Invalid regex %r: %s", pat, exc)
            continue

        if not matches:
            continue

        target = matches[match_index] if match_index < len(matches) else matches[-1]
        exists, result = _try_group(target, capture_group)
        if not exists or result is None:
            return None
        result = _apply_normalizers(result, normalize)
        if result is None:
            return None
        result = result.strip(strip_chars) if strip_chars else result.strip()
        return result or None

    return None


# ---------------------------------------------------------------------------
# Regex-array resolver  (all matches of a pattern → list of strings)
# ---------------------------------------------------------------------------

def _resolve_regex_all(rule: Dict[str, Any], text: str) -> List[Optional[str]]:
    """
    Like _resolve_regex but returns ALL matches as a list instead of one.

    Extra rule keys versus the scalar regex rule:
        max_items   int   Cap the number of results (default: unlimited).
    """
    primary = rule.get("pattern", "")
    if not primary:
        logger.warning("Regex-array rule missing 'pattern': %s", rule)
        return []

    flags         = _build_flags(rule)
    capture_group = rule.get("capture_group", 0)
    normalize     = rule.get("normalize", "")
    strip_chars   = rule.get("strip_chars", "")
    max_items     = rule.get("max_items")

    try:
        matches = list(re.finditer(primary, text, flags))
    except re.error as exc:
        logger.error("Invalid regex %r: %s", primary, exc)
        return []

    results: List[Optional[str]] = []
    for m in matches:
        exists, result = _try_group(m, capture_group)
        if not exists or result is None:
            continue
        result = _apply_normalizers(result, normalize)
        if result is None:
            continue
        result = result.strip(strip_chars) if strip_chars else result.strip()
        if result:
            results.append(result)
        if max_items is not None and len(results) >= max_items:
            break

    return results


# ---------------------------------------------------------------------------
# Entity resolver
# ---------------------------------------------------------------------------

def _resolve_entity(rule: Dict[str, Any], doc: Any) -> Optional[str]:
    if doc is None:
        logger.warning("Entity resolver received None doc — skipping")
        return None

    labels = rule.get("label")
    if isinstance(labels, str):
        labels = [labels]
    label_set = set(labels or [])

    match_index   = rule.get("match_index", 0)
    min_length    = rule.get("min_length", 1)
    exclude_pat   = rule.get("exclude_pattern", "")
    exclude_flags = _build_flags({"flags": rule.get("exclude_flags", [])})
    normalize     = rule.get("normalize", "")

    candidates = [
        ent.text for ent in doc.ents
        if ent.label_ in label_set
        and len(ent.text) >= min_length
        and not (exclude_pat and re.search(exclude_pat, ent.text, exclude_flags))
    ]

    if not candidates:
        return None

    result = candidates[match_index] if match_index < len(candidates) else candidates[-1]
    return _apply_normalizers(result, normalize)


# ---------------------------------------------------------------------------
# Entity-array resolver  (all matching entities → list)
# ---------------------------------------------------------------------------

def _resolve_entity_all(rule: Dict[str, Any], doc: Any) -> List[Optional[str]]:
    """
    Returns ALL entities matching the label filter as a list.

    Extra rule key:
        max_items   int   Cap the number of results (default: unlimited).
        unique      bool  Deduplicate while preserving order (default: False).
    """
    if doc is None:
        logger.warning("Entity-array resolver received None doc — skipping")
        return []

    labels = rule.get("label")
    if isinstance(labels, str):
        labels = [labels]
    label_set = set(labels or [])

    min_length    = rule.get("min_length", 1)
    exclude_pat   = rule.get("exclude_pattern", "")
    exclude_flags = _build_flags({"flags": rule.get("exclude_flags", [])})
    normalize     = rule.get("normalize", "")
    max_items     = rule.get("max_items")
    unique        = rule.get("unique", False)

    results: List[str] = []
    seen: set = set()

    for ent in doc.ents:
        if ent.label_ not in label_set:
            continue
        if len(ent.text) < min_length:
            continue
        if exclude_pat and re.search(exclude_pat, ent.text, exclude_flags):
            continue

        value = _apply_normalizers(ent.text, normalize)
        if not value:
            continue
        if unique:
            if value in seen:
                continue
            seen.add(value)

        results.append(value)
        if max_items is not None and len(results) >= max_items:
            break

    return results


# ---------------------------------------------------------------------------
# Token-attribute resolver
# ---------------------------------------------------------------------------

def _resolve_token_attr(rule: Dict[str, Any], doc: Any) -> Optional[str]:
    if doc is None:
        logger.warning("Token-attr resolver received None doc — skipping")
        return None

    attr        = rule.get("attr", "")
    match_index = rule.get("match_index", 0)
    normalize   = rule.get("normalize", "")

    candidates = [t.text for t in doc if getattr(t, attr, False)]
    if not candidates:
        return None

    result = candidates[match_index] if match_index < len(candidates) else candidates[-1]
    return _apply_normalizers(result, normalize)


# ---------------------------------------------------------------------------
# Built-in resolver registration
# ---------------------------------------------------------------------------

register_resolver("regex",        lambda rule, doc, text: _resolve_regex(rule, text))
register_resolver("entity",       lambda rule, doc, text: _resolve_entity(rule, doc))
register_resolver("token_attr",   lambda rule, doc, text: _resolve_token_attr(rule, doc))
# Array-producing leaf resolvers (used internally by the array node path):
register_resolver("regex_all",    lambda rule, doc, text: _resolve_regex_all(rule, text))
register_resolver("entity_all",   lambda rule, doc, text: _resolve_entity_all(rule, doc))


# ---------------------------------------------------------------------------
# Scalar field dispatcher  (returns str | None)
# ---------------------------------------------------------------------------

def _resolve_scalar_field(
    rule: Dict[str, Any],
    doc: Any,
    text: str,
) -> Optional[str]:
    src = rule.get("source_type")
    fn = _RESOLVERS.get(src)
    if fn is None:
        logger.warning("Unknown source_type %r — no resolver registered", src)
        return None
    return fn(rule, doc, text)


# ---------------------------------------------------------------------------
# Generic nested schema resolver
# ---------------------------------------------------------------------------
#
# Schema node shapes
# ──────────────────
#
# 1. LEAF  (scalar string)
#    {
#      "source_type": "regex" | "entity" | "token_attr" | <custom>,
#      ...resolver-specific keys...
#    }
#
# 2. OBJECT  (nested dict of named fields)
#    {
#      "type": "object",
#      "fields": {
#        "field_a": <schema_node>,
#        "field_b": <schema_node>,
#        ...
#      }
#    }
#
# 3. ARRAY  (repeated items)
#    {
#      "type": "array",
#
#      # --- how to split the text into per-item segments (optional) ---
#      # If omitted the whole text is the single segment (useful when
#      # the item schema itself fans out via regex_all / entity_all).
#      "split_pattern": "<regex>",   # splits text; each piece → one item
#      "split_flags":   ["DOTALL"],  # re flags for split_pattern
#
#      # --- what each item looks like ---
#      "items": <schema_node>
#      # Can be a leaf, an object, or even another array (any depth).
#    }
#
# Results
# ───────
#   LEAF   → str | None
#   OBJECT → {field: result, ...}   (all keys always present, value may be None)
#   ARRAY  → [result, ...]          (may be empty; each element mirrors item schema)


def _resolve_node(node: SchemaNode, doc: Any, text: str) -> ResultNode:
    """
    Recursively resolve a schema node against `text` / `doc`.
    Dispatches on node["type"] or falls back to scalar leaf resolution.
    """
    node_type = node.get("type")

    if node_type == "object":
        return _resolve_object_node(node, doc, text)

    if node_type == "array":
        return _resolve_array_node(node, doc, text)

    # No "type" key → treat as a scalar leaf rule
    return _resolve_scalar_field(node, doc, text)


def _resolve_object_node(
    node: SchemaNode,
    doc: Any,
    text: str,
) -> Dict[str, ResultNode]:
    """
    Resolve every field in node["fields"] and return a dict.
    Each field may itself be a leaf, object, or array — fully recursive.
    """
    fields: Dict[str, SchemaNode] = node.get("fields", {})
    result: Dict[str, ResultNode] = {}

    for field_name, child_node in fields.items():
        try:
            result[field_name] = _resolve_node(child_node, doc, text)
        except Exception as exc:
            logger.error(
                "Object field %r raised unexpectedly: %s", field_name, exc,
                exc_info=True,
            )
            result[field_name] = None

    return result


def _resolve_array_node(
    node: SchemaNode,
    doc: Any,
    text: str,
) -> List[ResultNode]:
    """
    Resolve an array node:

    Two operating modes, selected by whether "split_pattern" is present:

    MODE A — split_pattern present
        Split the text into N segments; resolve item schema against each
        segment with its own spaCy doc.  Good for table rows, repeated
        blocks, delimited records, etc.

    MODE B — no split_pattern
        Resolve item schema against the full text once.
        If the item schema is a leaf with source_type in {regex_all,
        entity_all} it returns a list natively.
        If the item schema returns a list → that IS the array.
        If it returns a scalar → wrap in [scalar].
        This handles "give me all ORG entities" without needing a split.
    """
    item_schema: SchemaNode  = node.get("items", {})
    split_pat: Optional[str] = node.get("split_pattern")
    max_items: Optional[int]  = node.get("max_items")

    results: List[ResultNode] = []

    if split_pat:
        # ── MODE A: segment-per-item ────────────────────────────────────
        try:
            flags    = _build_flags({"flags": node.get("split_flags", [])})
            segments = re.split(split_pat, text, flags=flags)
        except re.error as exc:
            logger.error("Invalid split_pattern %r: %s", split_pat, exc)
            return []

        nlp = _get_nlp()

        try:
            segment_docs = list(nlp.pipe(segments))
        except Exception as exc:
            logger.error("spaCy pipe failed on array segments: %s", exc, exc_info=True)
            segment_docs = [None] * len(segments)

        for seg_doc, seg_text in zip(segment_docs, segments):
            if not seg_text.strip():
                continue
            try:
                item_result = _resolve_node(item_schema, seg_doc, seg_text)
            except Exception as exc:
                logger.error(
                    "Array item resolve raised: %s", exc, exc_info=True
                )
                item_result = None

            results.append(item_result)
            if max_items is not None and len(results) >= max_items:
                break

    else:
        # ── MODE B: whole-text, native fan-out ─────────────────────────
        try:
            raw = _resolve_node(item_schema, doc, text)
        except Exception as exc:
            logger.error("Array item resolve raised: %s", exc, exc_info=True)
            return []

        if isinstance(raw, list):
            results = raw
        elif raw is not None:
            results = [raw]

        if max_items is not None:
            results = results[:max_items]

    return results


# ---------------------------------------------------------------------------
# Safe wrappers
# ---------------------------------------------------------------------------

def _safe_resolve_node(
    path: str,
    node: SchemaNode,
    doc: Any,
    text: str,
) -> ResultNode:
    """Resolve a node; isolate crashes so siblings still complete."""
    try:
        return _resolve_node(node, doc, text)
    except Exception as exc:
        logger.error(
            "Schema path %r raised unexpectedly: %s", path, exc, exc_info=True
        )
        return None


def _doc_for_text(nlp: Any, text: str) -> Any:
    """Run *text* through spaCy, returning None on failure instead of raising."""
    try:
        return next(iter(nlp.pipe([text])))
    except Exception as exc:
        logger.error("spaCy pipe failed: %s", exc, exc_info=True)
        return None


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def extract_fields(
    text: str,
    fields: Dict[str, SchemaNode],
) -> Dict[str, ResultNode]:
    """
    Extract fields from a single text string.

    `fields` is a flat dict of {name: schema_node}.  Each schema node may be
    a scalar leaf, an object node, or an array node — nested to any depth.

    Returns {field_name: result_or_None}.

    Backward-compatible: callers that pass flat scalar rules unchanged still work.
    """
    nlp = _get_nlp()
    doc = _doc_for_text(nlp, text)
    return {
        field: _safe_resolve_node(field, node, doc, text)
        for field, node in fields.items()
    }


def extract_schema(
    text: str,
    schema: SchemaNode,
) -> ResultNode:
    """
    Resolve a *single* schema node (which may be a leaf, object, or array)
    against `text`.

    Useful when the top-level result should itself be a list or a structured
    object rather than a flat dict of fields.

    Example
    -------
    schema = {
        "type": "array",
        "split_pattern": r"\\n\\n+",
        "items": {
            "type": "object",
            "fields": {
                "date":   {"source_type": "entity", "label": "DATE"},
                "amount": {"source_type": "regex",  "pattern": r"\\$[\\d,]+"},
            }
        }
    }
    result = extract_schema(invoice_text, schema)
    # → [{"date": "Jan 2024", "amount": "$1,200"}, ...]
    """
    nlp = _get_nlp()
    doc = _doc_for_text(nlp, text)
    return _safe_resolve_node("<root>", schema, doc, text)


def extract_fields_batch(
    texts: List[str],
    fields: Dict[str, SchemaNode],
) -> List[Dict[str, ResultNode]]:
    """
    Extract fields from a list of texts in a single top-level spaCy pipe pass.

    Returns one result dict per input text, in the same order.

    Note: array nodes with split_pattern trigger their own inner pipe call per
    text; the outer pass still processes the top-level texts efficiently.
    """
    nlp = _get_nlp()
    cleaned = [clean_text(t) for t in texts]

    try:
        docs = list(nlp.pipe(cleaned))
    except Exception as exc:
        logger.error("spaCy pipe failed: %s", exc, exc_info=True)
        docs = [None] * len(cleaned)

    return [
        {
            field: _safe_resolve_node(field, node, doc, text)
            for field, node in fields.items()
        }
        for doc, text in zip(docs, cleaned)
    ]


def extract_schema_batch(
    texts: List[str],
    schema: SchemaNode,
) -> List[ResultNode]:
    """
    Like extract_schema but processes a list of texts efficiently.

    Returns one ResultNode per input text, in the same order.
    """
    nlp = _get_nlp()
    cleaned = [clean_text(t) for t in texts]

    try:
        docs = list(nlp.pipe(cleaned))
    except Exception as exc:
        logger.error("spaCy pipe failed: %s", exc, exc_info=True)
        docs = [None] * len(cleaned)

    return [
        _safe_resolve_node(f"<root>[{i}]", schema, doc, text)
        for i, (doc, text) in enumerate(zip(docs, cleaned))
    ]