"""agents.py — Multi-Agent Supervisor -> Scraper -> Validator using Mistral AI."""
import os
from langchain_mistralai import ChatMistralAI
from langchain_groq import ChatGroq
from langgraph.prebuilt import create_react_agent
from langgraph_supervisor import create_supervisor
from langgraph.checkpoint.memory import MemorySaver

from tools import (
    search_openalex, search_tavily, search_scopus, search_apify_scholar,
    validate_papers, run_bertopic, upload_to_storage, classify_paper_types
)
from prompts import (
    RINGMASTER_SUPERVISOR_PROMPT,
    SCRAPER_AGENT_PROMPT,
    VALIDATOR_AGENT_PROMPT,
)

def build_agent():
    """Build the Multi-Agent graph."""
    
    # ── LLM Configuration w/ Fallbacks ──
    mistral_llm = ChatMistralAI(
        model="mistral-small-latest",
        api_key=os.getenv("MISTRAL_API_KEY"),
        temperature=0,
        max_tokens=512,
        max_retries=1
    )
    groq_llm = ChatGroq(
        model="llama-3.3-70b-versatile",
        api_key=os.getenv("GROQ_API_KEY"),
        temperature=0,
        max_tokens=512
    )
    llm = mistral_llm.with_fallbacks([groq_llm])

    # ── 1. Scraper Agent ──
    scraper_agent = create_react_agent(
        model=llm,
        tools=[search_openalex, search_tavily, search_scopus, search_apify_scholar],
        name="scraper_agent",
        prompt=SCRAPER_AGENT_PROMPT
    )

    # ── 2. Validator & Analysis Agent ──
    validator_agent = create_react_agent(
        model=llm,
        tools=[validate_papers, run_bertopic, classify_paper_types, upload_to_storage],
        name="validator_agent",
        prompt=VALIDATOR_AGENT_PROMPT
    )

    # ── 3. Supervisor Ringmaster ──
    workflow = create_supervisor(
        [scraper_agent, validator_agent],
        model=llm,
        prompt=RINGMASTER_SUPERVISOR_PROMPT,
        output_mode="full_history"
    )

    return workflow.compile(checkpointer=MemorySaver())