import logging
import os
import platform

from mm_agents.vlaa_gui.agents.grounding import ACI
from mm_agents.vlaa_gui.agents.worker import Worker
from mm_agents.vlaa_gui.core.engine import (
    OpenAIEmbeddingEngine,
    GeminiEmbeddingEngine,
    AzureOpenAIEmbeddingEngine,
    ArkEmbeddingEngine,
    QwenEmbeddingEngine,
)

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


class UIAgent:
    """Base class for UI automation agents"""

    def __init__(
        self,
        engine_params: dict,
        grounding_agent: ACI,
        platform: str = platform.system().lower(),
        action_space: str = "pyautogui",
        observation_type: str = "screenshot",
        search_engine: str = "llm",
    ):
        """Initialize UIAgent

        Args:
            engine_params: Configuration parameters for the LLM engine
            grounding_agent: Instance of ACI class for UI interaction
            platform: Operating system platform (macos, linux, windows)
            action_space: Type of action space to use (pyautogui, aci)
            observation_type: Type of observations to use (a11y_tree, mixed)
            engine: Search engine to use (perplexica, LLM)
        """
        self.engine_params = engine_params
        self.grounding_agent = grounding_agent
        self.platform = platform
        self.action_space = action_space
        self.observation_type = observation_type
        self.engine = search_engine

    def reset(self) -> None:
        """Reset agent state"""
        pass

    def predict(self, instruction: str, observation: dict) -> tuple[dict, list[str]]:
        """Generate next action prediction

        Args:
            instruction: Natural language instruction
            observation: Current UI state observation

        Returns:
            Tuple containing agent info dictionary and list of actions
        """
        pass


class Agent(UIAgent):
    """Agent that uses iterative planning for UI automation."""

    def __init__(
        self,
        engine_params: dict,
        grounding_agent: ACI,
        platform: str = platform.system().lower(),
        with_reflection: bool = True,
        action_space: str = "pyautogui",
        observation_type: str = "screenshot",
        search_engine: str | None = None,
        reflection_engine_params: dict | None = None,
        memory_root_path: str = os.getcwd(),
        memory_folder_name: str = "agent_memory",
        memory_type: str = "null",
        kb_release_tag: str = "v0.2.2",
        embedding_engine_type: str = "openai",
        embedding_engine_params: dict = {},
        coding_agent_flag: bool = False,
        max_image_in_trajectory: int = 8,
        debug: bool = False,
        lexical_weight: float = 0.0,
        enable_gate: bool = True,
        loop_detection: bool = True,
        feasibility_check: bool = True,
    ):
        """Initialize Agent

        Args:
            engine_params: Configuration parameters for the LLM engine
            grounding_agent: Instance of ACI class for UI interaction
            platform: Operating system platform (darwin, linux, windows)
            action_space: Type of action space to use (pyautogui, other)
            observation_type: Type of observations to use (a11y_tree, screenshot, mixed)
            search_engine: Search engine to use (LLM, perplexica)
            memory_root_path: Path to memory directory. Defaults to current working directory.
            memory_folder_name: Name of memory folder. Defaults to "agent_memory".
            kb_release_tag: Release tag for knowledge base. Defaults to "v0.2.2".
            embedding_engine_type: Embedding engine to use for knowledge base. Defaults to "openai". Supports "openai" and "gemini".
            embedding_engine_params: Parameters for embedding engine. Defaults to {}.
            debug: bool = False,
            max_image_in_trajectory: int = 8,
        """
        super().__init__(
            engine_params,
            grounding_agent,
            platform,
            action_space,
            observation_type,
            search_engine,
        )

        self.with_reflection = with_reflection
        self.reflection_engine_params = reflection_engine_params or engine_params
        self.memory_root_path = memory_root_path
        self.memory_folder_name = memory_folder_name
        self.lexical_weight = lexical_weight
        self.kb_release_tag = kb_release_tag
        self.memory_type = memory_type
        self.local_kb_path = os.path.join(
            self.memory_root_path, self.memory_folder_name
        )

        self.coding_agent_flag = coding_agent_flag
        self.max_image_in_trajectory = max_image_in_trajectory
        self.enable_gate = enable_gate
        self.loop_detection = loop_detection
        self.feasibility_check = feasibility_check

        if embedding_engine_type == "openai":
            self.embedding_engine = OpenAIEmbeddingEngine(**embedding_engine_params)
        elif embedding_engine_type == "gemini":
            self.embedding_engine = GeminiEmbeddingEngine(**embedding_engine_params)
        elif embedding_engine_type == "azure":
            self.embedding_engine = AzureOpenAIEmbeddingEngine(
                **embedding_engine_params
            )
        elif embedding_engine_type == "ark":
            self.embedding_engine = ArkEmbeddingEngine(**embedding_engine_params)
        elif embedding_engine_type == "qwen":
            self.embedding_engine = QwenEmbeddingEngine(**embedding_engine_params)
        else:
            raise ValueError(
                f"Unsupported embedding engine type: {embedding_engine_type}"
            )
        self.debug = debug
        self.reset()

    def reset(self) -> None:
        """Reset agent state and initialize components"""
        self.executor = Worker(
            engine_params=self.engine_params,
            grounding_agent=self.grounding_agent,
            reflection_engine_params=self.reflection_engine_params,
            local_kb_path=self.local_kb_path,
            embedding_engine=self.embedding_engine,
            platform=self.platform,
            enable_reflection=self.with_reflection,
            observation_type=self.observation_type,
            planning_type="iterative",
            search_engine=self.engine,
            memory_type=self.memory_type,
            use_task_experience=self.memory_type == "mixed",
            coding_agent_flag=self.coding_agent_flag,
            max_image_in_trajectory=self.max_image_in_trajectory,
            debug=self.debug,
            lexical_weight=self.lexical_weight,
            enable_gate=self.enable_gate,
            loop_detection=self.loop_detection,
            feasibility_check=self.feasibility_check,
        )

        # Reset state variables
        self.step_count: int = 0
        self.turn_count: int = 0
        self.search_query: str = ""
        self.verifier_feedback: dict | None = None

    def reset_executor_state(self) -> None:
        """Reset executor and step counter"""
        self.executor.reset()
        self.step_count = 0

    def restart_after_failed_verification(
        self, reason: str = "", missing_steps: str = ""
    ) -> None:
        """Force a fresh plan when the verifier rejects completion."""
        logger.info("Verifier rejected DONE. Forcing replan from scratch.")
        self.verifier_feedback = {
            "reason": reason,
            "missing_steps": missing_steps,
        }
        self.search_query = ""
        self.reset_executor_state()
        self.executor.set_replan_feedback(self.verifier_feedback)

    def predict(self, instruction: str, observation: dict) -> tuple[dict, list[str]]:
        result = self.executor.generate_next_action(observation, instruction)
        if self.verifier_feedback is not None and not self.executor.pending_verifier_feedback:
            self.verifier_feedback = None
        return result
