from google.adk import Agent, Runner
from google.adk.sessions import InMemorySessionService
from tools import create_plant_diagnosis_tool, create_remedy_retrieval_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.

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

            agent = Agent(
                name="AuraMindGlowAgent",
                model="gemini-2.0-flash",
                description="A farming assistant that can diagnose plant health and suggest remedies.",
                instruction="You are a friendly farming assistant. Your goal is to help users identify plant health issues and find solutions. Use your tools to diagnose the plant from an image and then find a remedy.",
                tools=[diagnosis_tool, remedy_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": remedy_tool,
                "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