app/services/background_tasks.py

"""Run long operations in a background daemon thread.

Streamlit reruns the whole script on every interaction (including clicking a
different step in the sidebar). A long, synchronous operation — generating a
proposal, running the analysis — gets interrupted by that rerun, so the work is
lost. This helper runs the operation on an independent daemon thread whose life
is NOT tied to the script run, keeps the result in a module-level registry, and
lets the UI poll for completion. The user can switch steps freely; the work
keeps going and the result is waiting when they return.

IMPORTANT: the worker function must not touch ``st.session_state`` (Streamlit
calls from a thread with no script-run context fail). Pass everything it needs
as arguments and read the result back on the main thread.

API-key propagation: the API key now lives only in the per-session store
(``st.session_state``), which a worker thread cannot see. So before spawning the
worker we snapshot the launching session's config overrides on the MAIN thread
and re-install them into the worker's own thread-local (see
``app.config.export_session_overrides`` / ``install_thread_overrides``). This
makes ``get_settings()`` inside the worker resolve the same key/provider/model
the user configured — without leaking it to any other visitor's worker. The LLM
client's usage logging is already guarded against a missing session.
"""

from __future__ import annotations

import threading
from typing import Any, Callable

from app import config

_lock = threading.Lock()
_tasks: dict[str, dict[str, Any]] = {}


def _has_streamlit_runtime() -> bool:
    """True only inside a real Streamlit script run (not tests / bare mode)."""
    try:
        from streamlit.runtime.scriptrunner import get_script_run_ctx

        return get_script_run_ctx() is not None
    except Exception:  # noqa: BLE001
        return False


def start(task_id: str, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> bool:
    """Start ``fn`` under ``task_id``.

    In the live app this runs on a daemon thread (so a Streamlit rerun can't
    kill it). Outside a real Streamlit runtime — unit tests, bare mode — it
    runs SYNCHRONOUSLY so callers that expect an immediate result still work.

    Returns True if a new run was started, False if one was already running
    (so a double-click or a rerun won't launch duplicate work).
    """
    with _lock:
        cur = _tasks.get(task_id)
        if cur and cur.get("status") == "running":
            return False
        _tasks[task_id] = {"status": "running", "result": None, "error": None}

    # Snapshot the launching session's config (incl. the session-only API key)
    # HERE, on the main thread, while st.session_state is reachable. The worker
    # re-installs it into its own thread-local so get_settings() resolves the
    # right key. Empty when no session (tests / bare mode).
    session_overrides = config.export_session_overrides()

    def _run() -> None:
        if session_overrides:
            config.install_thread_overrides(session_overrides)
        try:
            res = fn(*args, **kwargs)
            with _lock:
                _tasks[task_id] = {"status": "done", "result": res, "error": None}
        except Exception as exc:  # noqa: BLE001 - surface as a clean error status
            with _lock:
                _tasks[task_id] = {"status": "error", "result": None, "error": str(exc)}
        finally:
            if session_overrides:
                config.clear_thread_overrides()

    if _has_streamlit_runtime():
        threading.Thread(target=_run, daemon=True).start()
    else:
        _run()  # synchronous fallback for tests / bare mode
    return True


def status(task_id: str) -> dict[str, Any]:
    """Return ``{"status", "result", "error"}`` for ``task_id`` (status idle if unknown)."""
    with _lock:
        t = _tasks.get(task_id)
        return dict(t) if t else {"status": "idle", "result": None, "error": None}


def is_running(task_id: str) -> bool:
    return status(task_id).get("status") == "running"


def pop(task_id: str) -> dict[str, Any]:
    """Read and REMOVE a finished task's record (done/error). Running tasks stay."""
    with _lock:
        t = _tasks.get(task_id)
        if t and t.get("status") in ("done", "error"):
            _tasks.pop(task_id, None)
        return dict(t) if t else {"status": "idle", "result": None, "error": None}


def clear(task_id: str) -> None:
    with _lock:
        _tasks.pop(task_id, None)