from llama_index.tools.google.search.base import GoogleSearchToolSpec
from llama_index.core.tools.tool_spec.load_and_search import LoadAndSearchToolSpec
from utils.logger import logger
from config.settings import load_config

config = load_config()
google_config = config.get("google", {})

API_KEY = google_config.get("google_api_key")
CSE_ID = google_config.get("google_cse_id")

google_tools = []
try:
    google_spec = GoogleSearchToolSpec(key=API_KEY, engine=CSE_ID, num=5)
    google_tools = LoadAndSearchToolSpec.from_defaults(
        google_spec.to_tool_list()[0]
    ).to_tool_list()
    logger.info(
        f"Loaded Google search tools: {[t.metadata.name for t in google_tools]}"
    )
except Exception as e:
    logger.error(f"Failed to initialize GoogleSearchToolSpec: {e}")

# Export to web_search_agent
web_search_tools_llama = google_tools

from llama_index.core.agent.workflow import FunctionAgent
from tools.web_search_tool import web_search_tools_llama
from utils.logger import logger


def build_web_search_agent(llm):
    logger.info("Initializing Web Search agent and tool...")
    if not web_search_tools_llama:
        raise ValueError("Web search tool is not loaded.")

    agent = FunctionAgent(
        name="WebSearchAgent",
        description="Performs live web searches and returns summarized results.",
        system_prompt=(
            "You are a web search specialist. Use the provided tools to search online "
            "and deliver concise findings."
        ),
        llm=llm,
        tools=web_search_tools_llama,
    )
    logger.info("Web Search agent initialized successfully.")
    return agent

class ClinicalCDSSWorkflow(Workflow):
    def __init__(self, llm, query_engine, **kwargs):
        super().__init__(**kwargs)
        # Directly use the tools or specialized agents
        self.guideline_tool = guideline_query_tool_factory(query_engine)
        self.survival_prediction_tool = survival_prediction_tool
        self.web_search_agent = build_web_search_agent(llm)
        self.image_analysis_tool = all_image_analysis_tool
        self.papers_search_agent = build_papers_search_agent(llm)