from google.adk import Agent, Runner
from google.adk.sessions import InMemorySessionService
from tools import create_plant_diagnosis_tool, create_chroma_db_search_tool, create_bigquery_search_tool

def initialize_adk(vision_model, processor, retriever):
    """
    Initializes the ADK agent, tools, and runner.

    Args:
        vision_model: The loaded vision model.
        processor: The model's processor.
        retriever: The RAG retriever (no longer used but kept for compatibility).

    Returns:
        A dictionary containing the ADK components, or None on error.
    """
    print("Initializing ADK Tools and Agent...")
    try:
        if vision_model and processor:
            diagnosis_tool = create_plant_diagnosis_tool(model=vision_model, processor=processor)
            chroma_db_tool = create_chroma_db_search_tool()
            bigquery_tool = create_bigquery_search_tool()

            agent = Agent(
                name="AuraMindGlowAgent",
                model="gemini-2.0-flash",
                description="A farming assistant that can diagnose plant health and suggest remedies from local or cloud knowledge bases.",
                instruction="You are a friendly and intelligent farming assistant. Your primary goal is to help users identify plant health issues and find solutions. You have several tools at your disposal:\n" \
                            "- Use the 'diagnose_plant' tool when the user provides an image.\n" \
                            "- After diagnosing, use the 'search_chroma_db' tool to find a remedy in the local knowledge base.\n" \
                            "- If the user asks for cloud-based information or structured data, use the 'search_bigquery' tool.\n" \
                            "Always prioritize the local knowledge base (ChromaDB) for remedies unless the user specifies otherwise.",
                tools=[diagnosis_tool, chroma_db_tool, bigquery_tool]
            )

            session_service = InMemorySessionService()
            runner = Runner(agent=agent, app_name="AuraMindGlow", session_service=session_service)
            
            print("✅ ADK Agent and Runner initialized successfully!")
            return {
                "agent": agent,
                "runner": runner,
                "diagnosis_tool": diagnosis_tool,
                "remedy_tool": chroma_db_tool,  # For compatibility with app.py
                "session_service": session_service
            }
        else:
            print("❌ Skipping ADK setup due to errors in model or RAG loading.")
            return None
    except Exception as e:
        print(f"❌ CRITICAL ERROR during ADK setup: {e}")
        return None