"""
ui/agents.py
------------
Heavy model / graph loading, done once per server process and shared by every
browser session.

Streamlit gave this to us for free via @st.cache_resource. Gradio has no
equivalent, and it serves requests from a thread pool, so a naive "load if
None" would let two simultaneous first-visitors each start a ~2 GB model load.
Hence the explicit double-checked locking below: the lock is held across the
load, and a second caller blocks and then sees the finished object.
"""

import threading

_agents_lock = threading.Lock()
_agents = None

_chatbot_lock = threading.Lock()
_chatbot_models = None


def load_agents():
    """Import the compiled LangGraph agents. Importing the search graph also
    loads the SPECTER reranker model at module import time (by design)."""
    global _agents
    if _agents is None:
        with _agents_lock:
            if _agents is None:
                from app.modules.intent.graph import graph as intent_graph
                from app.modules.search.graph import graph as search_graph
                _agents = (intent_graph, search_graph)
    return _agents


def warm_chatbot_models():
    """Pre-load the chatbot's embedding + cross-encoder models so the first
    'Chat it out' click doesn't pay the model-load cost. We instantiate the
    exact models the chatbot uses (BAAI/bge-base-en-v1.5 + BAAI/bge-reranker-base),
    warming the weights into the HF/torch cache.

    Both are pinned to CPU: this runs at boot, outside any ZeroGPU window, and
    its whole job is to pull weights down — the GPU copy is made later, inside
    ui.gpu.vectorize_on_gpu, from the same warmed cache."""
    global _chatbot_models
    if _chatbot_models is None:
        with _chatbot_lock:
            if _chatbot_models is None:
                from vectorizeer import get_embeddings
                from langchain_community.cross_encoders import HuggingFaceCrossEncoder
                embeddings = get_embeddings(device="cpu")
                reranker = HuggingFaceCrossEncoder(
                    model_name="BAAI/bge-reranker-base",
                    model_kwargs={"device": "cpu"},
                )
                _chatbot_models = (embeddings, reranker)
    return _chatbot_models


def warm_chatbot_models_async():
    """Kick the chatbot warm-up onto a daemon thread.

    The Streamlit app warmed BOTH model sets behind one blocking splash, which
    meant nobody saw a usable page until ~2 GB of weights had downloaded. Only
    the agents are needed to act on the very first click, so we block on those
    and let the chatbot models finish in the background — they have until the
    user has framed an intent, run a search, and picked a paper, which is far
    longer than the load takes. ensure_chat_ready() calls warm_chatbot_models()
    anyway, so if the thread hasn't finished it simply blocks on the same lock.
    """
    threading.Thread(target=warm_chatbot_models, daemon=True).start()