"""Lightweight entity extraction for Phase 1 routing."""

from __future__ import annotations

import re

TICKER_AFTER_LABEL = re.compile(
    r"\bticker\s+([A-Z]{1,5})\b",
    re.IGNORECASE,
)
CAPS_TOKEN = re.compile(r"\b[A-Z]{1,5}\b")


def extract_entities(message: str) -> dict[str, str]:
    entities: dict[str, str] = {}

    label_match = TICKER_AFTER_LABEL.search(message)
    if label_match:
        entities["ticker"] = label_match.group(1).upper()
        return entities

    for token in CAPS_TOKEN.findall(message):
        if token not in {"I", "A", "AN", "THE", "AND", "OR", "FOR", "ON", "TO"}:
            entities["ticker"] = token
            break

    return entities