workers.py

import os
import signal
import time
from threading import Thread, Lock, Event

from logger import logger
from models import StockSources, MacroSource, NEWS, STOCK
from db_helpers import (
    read_sources,
    update_last_fetched,
    record_failure,
    store_raw_html,
    get_sources_due_for_market,
    update_source_table,
)
from config import (
    SCRAPER_TICK_SECONDS,
    SCRAPE_INTERVAL_MS,
    MAX_CONSECUTIVE_FAILURES,
    WORKER_ERROR_BACKOFF_SECONDS,
    MARKET_CONFIG,
)
from utils import (
    fetch_html,
    seconds_until_next_market_close,
    market_already_closed_today,
)


# =========================
# WORKER LIFECYCLE STATE
# =========================

workers_lock = Lock()

# Maps worker name → {"thread": Thread, "event": Event}
_workers: dict[str, dict] = {}


# =========================
# WORKER LOOP HARNESS
# =========================
# Each worker is split into a public function (entered by the Thread) and a
# private *_iteration that runs exactly one pass. The harness loops the
# iteration, catching unhandled exceptions so a surprise error in any one
# pass logs + backs off + continues, rather than killing the thread silently.


def _safe_worker_loop(name: str, iteration_fn, stop_event: Event):

    logger.info(f"{name} loop started")

    while not stop_event.is_set():
        try:
            iteration_fn(stop_event)
        except Exception:
            logger.exception(f"{name} crashed")
            stop_event.wait(WORKER_ERROR_BACKOFF_SECONDS)

    logger.info(f"{name} loop exiting (shutdown signaled)")


# =========================
# SCRAPER WORKER (news)
# =========================


def _scraping_iteration(stop_event: Event):
    """Scrape news sources and store their RAW HTML into the database."""

    logger.info("scraper tick: reading source list")
    sources = read_sources("news")
    now_ms = int(time.time() * 1000)

    logger.info(f"scraper tick: evaluating {len(sources)} sources")

    skipped_disabled = 0
    skipped_recent = 0
    attempted = 0

    for source in sources:

        if stop_event.is_set():
            logger.info("scraper tick: shutdown signaled mid-loop, returning")
            return

        if (source.failure_count or 0) >= MAX_CONSECUTIVE_FAILURES:
            logger.info(
                f"scraper: skipping disabled source {source.name} "
                f"(failures={source.failure_count})"
            )
            skipped_disabled += 1
            continue

        if (
            source.last_attempted_at is not None
            and now_ms - source.last_attempted_at < SCRAPE_INTERVAL_MS
        ):
            age_min = (now_ms - source.last_attempted_at) // 60000
            logger.info(
                f"scraper: skipping {source.name}, last attempt {age_min}min ago "
                f"(interval {SCRAPE_INTERVAL_MS // 60000}min)"
            )
            skipped_recent += 1
            continue

        source_name = source.name
        url = source.website
        attempted += 1

        logger.info(f"scraper: fetching {source_name} ({url})")

        try:
            html, error_reason = fetch_html(url)

            if html is None:
                logger.error(f"{source_name} fetch failed")
                logger.error(f"due to {error_reason}")
                record_failure(source.id, str(error_reason))
                continue

            raw_id = store_raw_html(source_name, html, NEWS)
            update_last_fetched(source.id)
            logger.info(f"Stored {source_name} -> {raw_id}")

        except Exception as e:
            logger.exception(f"{source_name} fetch failed")
            record_failure(source.id, str(e))

    logger.info(
        f"scraper tick done: attempted={attempted} "
        f"skipped_recent={skipped_recent} "
        f"skipped_disabled={skipped_disabled}; "
        f"sleeping {SCRAPER_TICK_SECONDS}s"
    )
    stop_event.wait(SCRAPER_TICK_SECONDS)


def scraping_worker(stop_event: Event):
    _safe_worker_loop("scraping_worker", _scraping_iteration, stop_event)


# =========================
# SHARED FETCH HELPER (stock + macro)
# =========================
# Both the stock and macro workers fetch a pre-built URL and store the raw HTML;
# only the source label and raw_html.type differ. update_source_table stamps
# last_fetched_at so the row drops out of "due" for the rest of the market day.


def _fetch_and_store(row, model, source_label: str, type_: str):
    """Fetch row.url and store it as raw_html of `type_`. Stamps the source row."""

    ident = getattr(row, "name", None) or getattr(row, "ticker", None) or row.id
    country = getattr(row, "country", None)
    url = row.url

    if not url:
        logger.error(f"{ident} has no pre-built url; skipping")
        update_source_table(model, row.id, "missing url")
        return

    logger.info(f"fetching {ident} ({type_})")

    try:
        html, error = fetch_html(url)

        if not html:
            update_source_table(model, row.id, error)
            logger.error(f"fetch failed for {ident}: {error}")
            return

        store_raw_html(source_label, html, type_, country=country)
        update_source_table(model, row.id)
        logger.info(f"{ident}: raw html stored ({type_})")

    except Exception:
        logger.exception(f"fetch failed for {ident}")
        update_source_table(model, row.id, "fetch exception")


def _sleep_until_next_close(stop_event: Event, label: str):
    """Sleep until the soonest market close across all markets (weekend-aware)."""

    sleep_s = min(seconds_until_next_market_close(m) for m in MARKET_CONFIG)
    logger.info(
        f"{label} tick done; sleeping {sleep_s / 3600:.1f}h until next market close"
    )
    stop_event.wait(sleep_s)


# =========================
# STOCK ENHANCER WORKER (fetch-only)
# =========================
# Once per market close (weekday), fetch every stock for that market that hasn't
# been fetched today, store its Google Finance HTML (type=stock), then sleep to
# the next close. All metric extraction happens on the LLM side.


def _stock_market_iteration(stop_event: Event):

    for market in MARKET_CONFIG:

        if not market_already_closed_today(market):
            logger.info(f"stock: {market} not closed yet today (or weekend); skipping")
            continue

        due = get_sources_due_for_market(StockSources, market)
        logger.info(f"stock: {market} closed; {len(due)} stock(s) due")

        for stock in due:
            if stop_event.is_set():
                logger.info("stock tick: shutdown signaled mid-loop, returning")
                return
            _fetch_and_store(stock, StockSources, "Google Finance", STOCK)

    _sleep_until_next_close(stop_event, "stock")


def stock_enrichment_worker(stop_event: Event):
    _safe_worker_loop("stock_enrichment_worker", _stock_market_iteration, stop_event)


# =========================
# MACRO WORKER (fetch-only)
# =========================
# Same once-per-close cadence as the stock worker, but over MacroSource rows
# (FII/DII, VIX, commodities, indices, FX, yields). The raw_html.type is taken
# from each row so one worker covers every macro family.


def _macro_iteration(stop_event: Event):

    for market in MARKET_CONFIG:

        if not market_already_closed_today(market):
            logger.info(f"macro: {market} not closed yet today (or weekend); skipping")
            continue

        due = get_sources_due_for_market(MacroSource, market)
        logger.info(f"macro: {market} closed; {len(due)} series due")

        for src in due:
            if stop_event.is_set():
                logger.info("macro tick: shutdown signaled mid-loop, returning")
                return
            _fetch_and_store(src, MacroSource, src.name, src.type)

    _sleep_until_next_close(stop_event, "macro")


def macro_worker(stop_event: Event):
    _safe_worker_loop("macro_worker", _macro_iteration, stop_event)


# =========================
# WORKER REGISTRY
# =========================
# Maps a stable name → the worker's entry-point function.
# Registered actions:
#   START_SCRAPER         / STOP_SCRAPER         / RESTART_SCRAPER
#   START_STOCK_ENHANCER  / STOP_STOCK_ENHANCER  / RESTART_STOCK_ENHANCER
#   START_MACRO           / STOP_MACRO           / RESTART_MACRO

_WORKER_REGISTRY: dict[str, callable] = {
    "scraper": scraping_worker,
    "stock_enhancer": stock_enrichment_worker,
    "macro": macro_worker,
}


# =========================
# INTERNAL HELPERS
# =========================


def _launch_worker(name: str):
    """Spin up a single named worker and register it in _workers.
    Caller must hold workers_lock."""

    if name not in _WORKER_REGISTRY:
        raise ValueError(
            f"Unknown worker: {name!r}. Valid names: {list(_WORKER_REGISTRY)}"
        )

    stop_event = Event()
    thread = Thread(
        target=_WORKER_REGISTRY[name],
        args=(stop_event,),
        name=name,
        daemon=True,
    )
    _workers[name] = {"thread": thread, "event": stop_event}
    thread.start()
    logger.info(f"_launch_worker: '{name}' started (tid={thread.ident})")


def _signal_and_join(name: str, stop_event: Event, thread: Thread, timeout: float):
    """Signal a worker's event and wait for its thread to exit."""

    logger.info(f"stop: signaling '{name}'")
    stop_event.set()
    thread.join(timeout=timeout)
    if thread.is_alive():
        logger.warning(f"stop: '{name}' did not exit within {timeout}s")
    else:
        logger.info(f"stop: '{name}' exited cleanly")


# =========================
# PUBLIC CONTROL API
# =========================


def start_workers(names: list[str] | None = None):
    """Start workers. If *names* is omitted, all registered workers are started.
    Already-running workers are skipped with a warning."""

    targets = names or list(_WORKER_REGISTRY.keys())

    with workers_lock:
        for name in targets:
            entry = _workers.get(name)
            if entry and entry["thread"].is_alive():
                logger.warning(f"start_workers: '{name}' is already running — skipping")
                continue
            _launch_worker(name)

    logger.info(f"start_workers: launched {targets}")


def stop_worker(name: str, timeout: float = 5.0):
    """Stop a single named worker and remove it from the registry."""

    with workers_lock:
        entry = _workers.get(name)
        if not entry:
            logger.warning(f"stop_worker: '{name}' is not running — nothing to do")
            return

        _signal_and_join(name, entry["event"], entry["thread"], timeout)
        del _workers[name]


def stop_workers(names: list[str] | None = None, timeout: float = 5.0):
    """Stop workers. If *names* is omitted, all running workers are stopped."""

    targets = names or list(_workers.keys())
    for name in targets:
        stop_worker(name, timeout=timeout)
    logger.info(f"stop_workers: done — {targets}")


def restart_worker(name: str, timeout: float = 5.0):
    """Stop a worker (if running) then immediately restart it."""

    with workers_lock:
        entry = _workers.get(name)
        if entry:
            _signal_and_join(name, entry["event"], entry["thread"], timeout)
            del _workers[name]
        _launch_worker(name)
    logger.info(f"restart_worker: '{name}' restarted")


def worker_status() -> dict[str, bool]:
    """Return a snapshot of {worker_name: is_alive} for every running worker."""

    with workers_lock:
        return {name: entry["thread"].is_alive() for name, entry in _workers.items()}


# =========================
# PROCESS SHUTDOWN
# =========================


def shutdown_server():
    logger.info("shutdown_server: stopping all workers then sending SIGTERM")
    stop_workers()
    # SIGTERM is the semantically correct "please exit" signal; friendlier to
    # process managers (systemd, supervisord, docker) than SIGINT.
    os.kill(os.getpid(), signal.SIGTERM)