stock-oracle/app/services/attention/entity_resolver.py

"""
Entity Resolver — maps ticker symbols to canonical company entities.

Resolution pipeline:
1. Look up company name from `companies` table
2. Normalize name (strip legal suffixes) → canonical_name + aliases
3. Validate against Wikipedia Search API → best wiki_title + confidence
4. Build GDELT query from canonical_name + aliases
5. Upsert into company_entity_map (respects is_manual_override flag)
"""

import logging
import re
from typing import Optional

import httpx
from sqlalchemy import select, update
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.attention import CompanyEntityMap
from app.models.financial import Company

logger = logging.getLogger(__name__)

# Legal suffixes to strip for canonical name derivation
_SUFFIX_PATTERN = re.compile(
    r",?\s+\b(Inc\.?|Corp\.?|Corporation|Holdings?|Ltd\.?|Limited|LLC|L\.L\.C\.|"
    r"Group|Co\.?|Company|Technologies|Technology|International|Industries|"
    r"Pharmaceuticals?|Therapeutics?|Sciences?|Bancorp|Financial|Holding|"
    r"Acquisition|Acquisitions|Capital|Partners|Trust|"
    r"Com)\s*$",  # "Com" catches SEC-style ".com" artifacts (e.g. "AMAZON COM")
    re.IGNORECASE,
)

_COMPANY_KEYWORDS = {
    "company", "corporation", "inc", "corp", "ltd", "llc", "holdings",
    "stock", "shares", "nasdaq", "nyse", "ticker", "finance", "financial",
    "investor", "business", "enterprise", "industries",
}


def _normalize_name(raw_name: str) -> tuple[str, list[str]]:
    """Return (canonical_name, aliases).

    Strips common legal suffixes iteratively to produce a short canonical name.
    Preserves the original and intermediate forms as aliases.
    """
    aliases = []
    current = raw_name.strip()

    for _ in range(5):  # max 5 iterations to avoid infinite loops
        stripped = _SUFFIX_PATTERN.sub("", current).strip().rstrip(",").strip()
        if stripped == current or not stripped:
            break
        aliases.append(current)
        current = stripped

    canonical = current
    # Also add the fully original name if not already captured
    if raw_name.strip() != canonical and raw_name.strip() not in aliases:
        aliases.insert(0, raw_name.strip())

    # Deduplicate while preserving order
    seen = set()
    unique_aliases = []
    for a in aliases:
        if a not in seen and a != canonical:
            seen.add(a)
            unique_aliases.append(a)

    return canonical, unique_aliases


_WIKI_HEADERS = {
    "User-Agent": "StockOracleBot/1.0 (attention-subsystem; contact@stockoracle.internal)"
}

_SEC_HEADERS = {
    "User-Agent": "Stock Oracle contact@stockoracle.internal",
    "Accept": "application/json",
}

# Module-level cache: populated on first SEC fetch, maps ticker.upper() → company title
_SEC_TICKER_MAP: dict[str, str] = {}


def _is_placeholder_name(ticker: str, canonical_name: str) -> bool:
    """Detect if the company name is just the ticker symbol repeated."""
    return canonical_name.upper() == ticker.upper()


async def _fetch_sec_company_name(ticker: str) -> Optional[str]:
    """Fetch real company name from SEC company_tickers.json.

    Caches the full ticker→name mapping on first call (~13K entries).
    Returns the SEC `title` field for the given ticker, or None if not found.
    """
    global _SEC_TICKER_MAP
    if not _SEC_TICKER_MAP:
        url = "https://www.sec.gov/files/company_tickers.json"
        try:
            async with httpx.AsyncClient(timeout=15.0, headers=_SEC_HEADERS) as client:
                resp = await client.get(url)
                resp.raise_for_status()
                data = resp.json()
            _SEC_TICKER_MAP = {
                v["ticker"].upper(): v["title"]
                for v in data.values()
                if "ticker" in v and "title" in v
            }
            logger.info("Loaded %d tickers from SEC company_tickers.json", len(_SEC_TICKER_MAP))
        except Exception as exc:
            logger.warning("Failed to fetch SEC company_tickers.json: %s", exc)
            return None

    title = _SEC_TICKER_MAP.get(ticker.upper())
    if title and title == title.upper() and not title.isnumeric():
        # SEC stores many names in ALL-CAPS (e.g. "AMAZON COM INC") — title-case them
        title = title.title()
    return title


async def _search_wikipedia(query: str) -> list[dict]:
    """Call Wikipedia Search API and return raw results list."""
    url = "https://en.wikipedia.org/w/api.php"
    params = {
        "action": "query",
        "list": "search",
        "srsearch": query,
        "srlimit": 5,
        "format": "json",
    }
    try:
        async with httpx.AsyncClient(timeout=10.0, headers=_WIKI_HEADERS) as client:
            resp = await client.get(url, params=params)
            resp.raise_for_status()
            data = resp.json()
            return data.get("query", {}).get("search", [])
    except Exception as exc:
        logger.warning("Wikipedia search failed for %r: %s", query, exc)
        return []


def _score_wiki_result(result: dict, canonical_name: str, aliases: list[str]) -> float:
    """Score a Wikipedia search result [0, 1] for company relevance.

    Scoring strategy:
    - Direct title match (title == canonical or alias): 0.9+
    - Title starts with canonical: 0.7+
    - Title contains canonical: 0.5 base
    - No canonical in title: 0.0
    - Non-company signals (album, film, etc.): capped at 0.1
    """
    raw_title = result.get("title", "")
    title = raw_title.lower()
    snippet = result.get("snippet", "").lower()
    combined = title + " " + snippet

    # Penalize obvious non-company pages immediately
    non_company_signals = ["album", "film", "television", "televisión", "song", "band",
                           "musician", "footballer", "athlete", "politician", "novel",
                           "book", "discography", "filmography", "radio station",
                           "broadcasting", "tv channel", "telev"]
    if any(sig in combined for sig in non_company_signals):
        return 0.1

    # Require snippet to have at least one finance keyword for non-exact-match titles
    finance_signals = ["company", "corporation", "stock", "nasdaq", "nyse", "shares",
                       "investor", "business", "enterprise", "holdings", "inc."]
    if not any(sig in combined for sig in finance_signals):
        return 0.15

    canonical_lower = canonical_name.lower()
    all_names = [canonical_name] + aliases
    all_names_lower = [n.lower() for n in all_names]

    # Check for exact title match (e.g. "Apple Inc." == alias "Apple Inc.")
    raw_title_stripped = raw_title.strip()
    for name in all_names:
        if raw_title_stripped.lower() == name.lower():
            return 0.95  # exact match

    # Check if title starts with canonical name
    title_starts_with_canonical = title.startswith(canonical_lower)
    title_contains_canonical = canonical_lower in title

    # Also check aliases
    title_starts_with_alias = any(title.startswith(n) for n in all_names_lower)
    title_contains_alias = any(n in title for n in all_names_lower)

    if not (title_contains_canonical or title_contains_alias):
        return 0.0

    # Check for pages that are ABOUT the company (vs. lists, histories, etc.)
    # Penalize titles that suggest subsidiary/list/history pages
    secondary_patterns = ["list of", "history of", "acquisition", "merger",
                          "subsidiary", "criticism", "controversy"]
    if any(pat in title for pat in secondary_patterns):
        return 0.25

    # Base score
    if title_starts_with_canonical or title_starts_with_alias:
        base = 0.7
    elif title_contains_canonical or title_contains_alias:
        base = 0.5
    else:
        base = 0.0

    # Reward company/finance keywords in title or snippet
    keyword_hits = sum(1 for kw in _COMPANY_KEYWORDS if kw in combined)
    keyword_score = min(keyword_hits / 3, 1.0)

    return min(base + keyword_score * 0.3, 1.0)


async def _resolve_wiki(
    canonical_name: str,
    aliases: list[str],
) -> tuple[Optional[str], float]:
    """Return (wiki_title, confidence). wiki_title is None if confidence < 0.5."""
    # Try multiple search strategies, prefer more specific queries first
    search_queries = []
    # Try full original name first (e.g. "Apple Inc.")
    for alias in aliases[:2]:
        search_queries.append(alias)
    # Then canonical
    search_queries.append(f"{canonical_name} company")

    best_title: Optional[str] = None
    best_score = 0.0

    for query in search_queries:
        results = await _search_wikipedia(query)
        for result in results:
            score = _score_wiki_result(result, canonical_name, aliases)
            if score > best_score:
                best_score = score
                best_title = result["title"]
        if best_score >= 0.85:
            break  # good enough — no need to try more queries

    if best_score < 0.5:
        return None, best_score

    return best_title, best_score


def _build_gdelt_query(canonical_name: str, aliases: list[str]) -> str:
    """Build a GDELT DOC API query string using quoted phrase OR logic."""
    terms = [canonical_name] + [a for a in aliases if a != canonical_name]
    quoted = [f'"{t}"' for t in terms[:4]]  # cap at 4 terms
    return " OR ".join(quoted)


async def resolve_entity(
    db: AsyncSession,
    ticker: str,
) -> CompanyEntityMap:
    """Resolve ticker → entity and upsert into company_entity_map.

    Returns the upserted CompanyEntityMap row.
    Raises ValueError if the ticker cannot be resolved.
    """
    ticker = ticker.upper()

    # Check if manual override already exists — skip re-resolution
    existing_result = await db.execute(
        select(CompanyEntityMap).where(CompanyEntityMap.ticker == ticker)
    )
    existing = existing_result.scalars().first()
    if existing and existing.is_manual_override:
        logger.info("Skipping resolution for %s: manual override active", ticker)
        return existing

    # Look up company name from companies table
    company_result = await db.execute(
        select(Company).where(Company.ticker == ticker)
    )
    company = company_result.scalars().first()
    if not company:
        raise ValueError(f"Ticker {ticker!r} not found in companies table")

    raw_name = company.name
    canonical_name, aliases = _normalize_name(raw_name)

    # Detect placeholder names like "AMZN Corporation" → canonical becomes "AMZN"
    if _is_placeholder_name(ticker, canonical_name):
        logger.info(
            "Detected placeholder name for %s (%r) — querying SEC for real name",
            ticker, raw_name,
        )
        sec_name = await _fetch_sec_company_name(ticker)
        if sec_name:
            logger.info("SEC fallback for %s: %r → %r", ticker, raw_name, sec_name)
            # Update the DB so next resolve skips this path
            await db.execute(
                update(Company).where(Company.ticker == ticker).values(name=sec_name)
            )
            canonical_name, aliases = _normalize_name(sec_name)
        else:
            logger.warning("SEC fallback found no name for %s, proceeding with placeholder", ticker)

    logger.info("Resolving entity for %s: canonical=%r aliases=%r", ticker, canonical_name, aliases)

    wiki_title, confidence = await _resolve_wiki(canonical_name, aliases)
    gdelt_query = _build_gdelt_query(canonical_name, aliases)

    logger.info(
        "Entity resolved: %s → wiki_title=%r confidence=%.2f",
        ticker, wiki_title, confidence,
    )

    # Upsert into company_entity_map
    stmt = (
        pg_insert(CompanyEntityMap)
        .values(
            ticker=ticker,
            canonical_name=canonical_name,
            wiki_title=wiki_title,
            gdelt_query=gdelt_query,
            aliases_json=aliases,
            resolver_confidence=confidence,
            is_manual_override=False,
        )
        .on_conflict_do_update(
            index_elements=["ticker"],
            set_=dict(
                canonical_name=canonical_name,
                wiki_title=wiki_title,
                gdelt_query=gdelt_query,
                aliases_json=aliases,
                resolver_confidence=confidence,
            ),
            where=CompanyEntityMap.is_manual_override == False,  # noqa: E712
        )
        .returning(CompanyEntityMap)
    )

    result = await db.execute(stmt)
    await db.commit()

    row = result.scalars().first()
    if row is None:
        # Manual override prevented update — return existing
        existing_result2 = await db.execute(
            select(CompanyEntityMap).where(CompanyEntityMap.ticker == ticker)
        )
        row = existing_result2.scalars().first()

    return row