import logging
from mm_agents.vlaa_gui.memory.procedural_memory import PROCEDURAL_MEMORY

import os


logger = logging.getLogger("desktopenv.searcher_agent")


# Agent action decorator
def searcher_agent_action(func):
    func.is_searcher_agent_action = True
    return func


# --- Abstract Base Class and Factory ---
class SearcherAgent:
    def __init__(self, engine_params: dict, platform: str):
        self.engine_params = engine_params
        self.result_dir = ""
        self.tutorial_or_hint = ""
        self.tutorial_notes = []
        self.max_trajectory_length = 8
        self.platform = platform
        self.budget = engine_params.get("budget", 20)

    @staticmethod
    def create(
        engine_params: dict,
        search_env,
        grounder_agent,
        platform: str,
        client_password: str = "password",
    ):
        searcher_type = engine_params.get("type", "vlm")
        # Fall back to LLM-based search if no search environment is available
        if searcher_type == "llm":
            return LLMSearcherAgent(
                engine_params=engine_params,
                platform=platform,
            )
        else:
            raise NotImplementedError(f"Unknown searcher type: {searcher_type}")

    def _get_search_time(self) -> int:
        """for the name of result directory"""
        if not self.result_dir:
            return 1
        search_times: list[int] = []
        try:
            if not os.path.exists(self.result_dir):
                return 1
            for item_name in os.listdir(self.result_dir):
                full_path = os.path.join(self.result_dir, item_name)
                if os.path.isdir(full_path) and item_name.startswith("search_"):
                    try:
                        time_val = int(item_name.split("_", 1)[1])
                        search_times.append(time_val)
                    except (ValueError, IndexError):
                        continue
        except Exception:
            return 1
        if not search_times:
            return 1
        return max(search_times) + 1

    def search(self, query: str, obs) -> str:
        """
        Args:
            query: Format like "How to xxxx?", must be a detailed subtask
            obs: Current screenshot
        """
        raise NotImplementedError("Subclasses must implement the 'search' method")


class LLMSearcherAgent(SearcherAgent):
    """
    LLM-based searcher agent that uses a language model with search capabilities
    to retrieve information without needing a visual environment.
    """

    def __init__(self, engine_params: dict, platform: str):
        SearcherAgent.__init__(self, engine_params=engine_params, platform=platform)

        provider = engine_params.get("engine_type", None)

        if provider == "openai":
            from mm_agents.vlaa_gui.core.engine import LMMEngineOpenAISearch

            self.search_engine = LMMEngineOpenAISearch(
                base_url=engine_params.get("base_url"),
                api_key=engine_params.get("api_key"),
                model=engine_params.get("model", "gpt-5.4"),
                rate_limit=engine_params.get("rate_limit", -1),
            )
        else:  # gemini (default) or genai
            from mm_agents.vlaa_gui.core.engine import LMMEngineGeminiSearch

            self.search_engine = LMMEngineGeminiSearch(
                api_key=engine_params.get("api_key"),
                api_keys=engine_params.get("api_keys"),
                model=engine_params.get("model", "gemini-3.1-pro-preview"),
                rate_limit=engine_params.get("rate_limit", -1),
            )

        logger.info(
            f"[LLMSearcher] Initialized with provider={provider}, model={engine_params.get('model')}"
        )

    def search(self, query: str, obs) -> str:
        """
        Perform a search using the LLM with search capabilities.

        Args:
            query: Format like "How to xxxx?", must be a detailed subtask
            obs: Current screenshot (not used in LLM search but kept for interface compatibility)

        Returns:
            dict: Search result containing tutorial or hint
        """
        logger.info(f"[LLMSearcher] Starting search for query: {query}")
        if PROCEDURAL_MEMORY._is_macos_platform(self.platform):
            logger.info("[LLMSearcher] Using macOS prompt context.")

        # Create a search prompt that instructs the LLM to find relevant information
        search_prompt = PROCEDURAL_MEMORY.construct_llm_searcher_prompt(
            query=query, platform=self.platform
        )
        system_message = PROCEDURAL_MEMORY.construct_llm_searcher_system_prompt(
            platform=self.platform
        )

        try:
            # Create messages in OpenAI format
            messages = [
                {
                    "role": "system",
                    "content": system_message,
                },
                {"role": "user", "content": search_prompt},
            ]

            # Generate response using the search-enabled engine
            response = self.search_engine.generate(
                messages=messages,
            )

            logger.info("[LLMSearcher] Successfully retrieved search results")

            # Return in the format expected by the searcher interface
            return {
                "query": query,
                "completion_reason": "DONE",
                "tutorial_notes": [response],
                "execution_history": [],
                "steps_executed": 1,
                "budget": self.budget,
                "final_answer": response,
            }

        except Exception as e:
            logger.error(f"[LLMSearcher] Search failed with error: {e}")
            error_message = f"Search failed: {str(e)}"

            return {
                "query": query,
                "completion_reason": "FAIL",
                "tutorial_notes": [],
                "execution_history": [],
                "steps_executed": 0,
                "budget": self.budget,
                "final_answer": error_message,
            }
