import os
import logging
from langchain_core.tools import tool
from langchain_groq import ChatGroq
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver

logger = logging.getLogger(__name__)

_SESSIONS = None
_CURRENT_THREAD_ID = None
_API_CLIENT = None
_RETRIEVER = None

def set_session_storage(sessions_dict: dict) -> None:
    global _SESSIONS
    _SESSIONS = sessions_dict

def set_current_thread_id(thread_id: str) -> None:
    global _CURRENT_THREAD_ID
    _CURRENT_THREAD_ID = thread_id

def set_api_client(api_client) -> None:
    global _API_CLIENT
    _API_CLIENT = api_client

SYSTEM_PROMPT = """You are DermaScan AI, a dermatology assistant.

LANGUAGE: reply in the same language the user used (Arabic or English).

STRICT RAG RULE: your knowledge is limited to:
1. knowledge-base via search_knowledge_base
2. AI image results from analysis
3. this conversation

Never invent symptoms, history, numbers. If info isn't available say: "This information is not available."

TOOLS:
- search_knowledge_base: for general clinical/knowledge questions (cite the source filename).
- Keep replies SHORT (2-4 sentences) unless the user asks for more detail.
"""

@tool
def search_knowledge_base(query: str) -> str:
    """Search the dermatology knowledge base for clinical information."""
    if _RETRIEVER is None:
        return "Retriever not initialized."
    docs = _RETRIEVER.invoke(query)
    if not docs:
        return "No relevant information found."
    formatted = []
    for i, d in enumerate(docs, 1):
        source = d.metadata.get("source", "unknown").replace("\\", "/").split("/")[-1]
        formatted.append(f"[Source {i}: {source}]\n{d.page_content}")
    return "\n\n---\n\n".join(formatted)

def build_agent(retriever, model_name: str = "llama-3.1-8b-instant", temperature: float = 0.25):
    global _RETRIEVER
    _RETRIEVER = retriever

    groq_key = os.environ.get("GROQ_API_KEY")
    if not groq_key:
        raise ValueError("GROQ_API_KEY not set in environment variables.")

    llm = ChatGroq(
        groq_api_key=groq_key,
        model_name=model_name,
        temperature=temperature,
        max_tokens=512,
        request_timeout=120,
    )

    tools = [search_knowledge_base]
    memory = MemorySaver()

    return create_react_agent(
        model=llm,
        tools=tools,
        prompt=SYSTEM_PROMPT,
        checkpointer=memory,
    )