"""PointerAgent — multi-agent orchestrator for OSWorld benchmark tasks."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from desktop_env.desktop_env import DesktopEnv

from .agent_executor import Executor
from .agent_feasibility_gate import FeasibilityGate
from .agent_planner import Planner
from .agent_verifier import Verifier
from .config import config
from .knowledge import KnowledgeBank
from .task_logger import get_agent_logger, save_screenshot, setup_task_logger
from .types import (
    AgentState,
    ExecutorAction,
    ExecutorDecision,
    Phase,
    VerifierDecision,
)
from .utils import APIProvider, resize_and_encode_png

_PROVIDER_MAP: dict[str, APIProvider] = {
    "aws": APIProvider.BEDROCK,
    "bedrock": APIProvider.BEDROCK,
    "anthropic": APIProvider.ANTHROPIC,
    "vertex": APIProvider.VERTEX,
}


class PointerAgent:
    """Multi-agent CUA for the OSWorld benchmark.

    Composes a planner, a consolidated executor (with direct GUI, coding, and
    file tools), and a verifier.  All agents communicate via the shared
    :class:`KnowledgeBank` and the persistent :class:`AgentState`.

    The harness calls ``predict(instruction, obs)`` in a loop.  Each call
    returns ``(response_text, actions)`` where *actions* is a list the
    harness feeds to ``env.step()``.
    """

    def __init__(
        self,
        env: DesktopEnv,
        provider_name: str = "aws",
        screen_size: tuple[int, int] = (1920, 1080),
        **kwargs: Any,
    ) -> None:
        self.env = env
        self.screen_size = screen_size
        self.provider = _PROVIDER_MAP.get(provider_name, APIProvider.BEDROCK)

        # Will be initialized in reset()
        self.knowledge: KnowledgeBank | None = None
        self.gate: FeasibilityGate | None = None
        self.planner: Planner | None = None
        self.executor: Executor | None = None
        self.verifier: Verifier | None = None
        self.state = AgentState()
        self.instruction: str = ""

    def reset(self, instruction: str, task_logger: logging.Logger, task_results_dir: str) -> None:
        """Reset all agent state for a new task."""
        setup_task_logger(task_logger, task_results_dir)
        self._task_log = get_agent_logger("orchestrator")
        self.instruction = instruction

        self.state = AgentState()

        self.knowledge = KnowledgeBank()
        self.knowledge.get_app_versions(self.env)

        self.gate = FeasibilityGate(
            env=self.env,
            knowledge=self.knowledge,
            screen_size=self.screen_size,
        )
        self.planner = Planner(env=self.env, knowledge=self.knowledge)
        self.executor = Executor(
            env=self.env,
            screen_size=self.screen_size,
        )
        self.verifier = Verifier(
            env=self.env,
        )

    def predict(
        self,
        obs: dict[str, Any] | None = None,
    ) -> tuple[str, list[Any]]:
        """Called by the harness in a loop.

        Returns ``(response_text, actions)`` where each action is either:
        - a pyautogui code string
        - ``"DONE"`` / ``"FAIL"`` / ``"WAIT"``
        - a dict with ``{"action_type": "DONE"/"FAIL"/"WAIT"}``

        State Machine:

        START -> FEASIBILITY_CHECK
        FEASIBILITY_CHECK -> PLANNING (feasible)
        FEASIBILITY_CHECK -> FAILED (infeasible)
        PLANNING -> EXECUTING
        EXECUTING -> GUI (gui action emitted)
        EXECUTING -> VERIFYING (executor says DONE or FAIL)
        GUI -> EXECUTING (next predict() call)
        VERIFYING -> DONE
        VERIFYING -> FAILED
        VERIFYING -> REPLANNING
        REPLANNING -> EXECUTING
        """
        screenshot_b64 = resize_and_encode_png(obs.get("screenshot", b"")) if obs else ""
        self.state.latest_screenshot_b64 = screenshot_b64
        self.state.step_count += 1

        self._task_log.info("\nStep %d — Phase: %s", self.state.step_count, self.state.phase.value)
        save_screenshot(screenshot_b64, self.state.step_count)

        # Pass-through phases (transition within same predict() call)
        if self.state.phase == Phase.FEASIBILITY_CHECK:
            self._handle_feasibility_check(screenshot_b64)

        if self.state.phase == Phase.GUI:
            self._handle_gui()

        if self.state.phase == Phase.PLANNING:
            self._handle_planning(screenshot_b64)

        # Action phases (each returns directly)
        if self.state.phase == Phase.EXECUTING:
            return self._handle_executing(screenshot_b64)

        if self.state.phase == Phase.VERIFYING:
            return self._handle_verification_sentinel()

        if self.state.phase == Phase.REPLANNING:
            return self._handle_replanning(screenshot_b64)

        if self.state.phase == Phase.DONE:
            return "Task completed.", [{"action_type": "DONE"}]

        if self.state.phase == Phase.FAILED:
            return "Task failed.", [{"action_type": "FAIL"}]

        return "Unknown state.", [{"action_type": "FAIL"}]

    # MARK: Phase Handlers

    def _handle_feasibility_check(self, screenshot_b64: str) -> None:
        result = self.gate.run(self.instruction, screenshot_b64)
        if result.infeasible:
            self.state.phase = Phase.FAILED
            self._task_log.info("\n=== GATE: INFEASIBLE ===\n%s", result.reason)
        else:
            self.state.phase = Phase.PLANNING
            self._task_log.info("\n=== GATE: FEASIBLE ===\n%s", result.reason)

    def _handle_gui(self) -> None:
        self._task_log.info("\n=== GUI -> EXECUTING ===")
        self.state.phase = Phase.EXECUTING

    def _handle_planning(self, screenshot_b64: str) -> None:
        plan = self.planner.run(self.instruction, screenshot_b64)
        self.state.plan = plan
        self.state.phase = Phase.EXECUTING
        self._task_log.info("\n=== PLANNING COMPLETE ===")
        self._task_log.info("Plan:\n%s", plan.format())

    def _handle_executing(self, screenshot_b64: str) -> tuple[str, list[Any]]:
        action, self.state = self.executor.run(self.instruction, self.state, screenshot_b64)

        if action is None:
            self._task_log.info("\n=== EXECUTION: LOOP LIMIT ===")
            return "Internal loop limit — waiting.", [{"action_type": "WAIT"}]

        if isinstance(action, ExecutorAction):
            self._task_log.info(
                "\n=== EXECUTION: %s ===\n%s",
                action.decision.value.upper(),
                action.reason,
            )
            self.state.executor_action = action
            self.state.phase = Phase.VERIFYING
            return "Verifying executor decision.", [{"action_type": "WAIT"}]

        # GUI action (str)
        self._task_log.info("\n=== EXECUTION STEP (step %d) ===", self.state.step_count)
        self.state.executor_action = ExecutorAction(ExecutorDecision.GUI, "Waiting for GUI action")
        self.state.phase = Phase.GUI
        return "Executing GUI action.", [action]

    def _handle_verification_sentinel(self) -> tuple[str, list[Any]]:
        """Map executor's terminal action directly to a final phase, skipping verification."""
        action = self.state.executor_action
        if action and action.decision == ExecutorDecision.FAIL:
            self.state.phase = Phase.FAILED
            self._task_log.info("\n=== SENTINEL: FAIL (no verification) ===\n%s", action.reason)
            return f"Executor declared failure: {action.reason}", [{"action_type": "FAIL"}]

        self.state.phase = Phase.DONE
        self._task_log.info(
            "\n=== SENTINEL: DONE (no verification) ===\n%s",
            action.reason if action else "(no action)",
        )
        return "Executor declared task complete.", [{"action_type": "DONE"}]

    def _handle_verifying(self, screenshot_b64: str) -> tuple[str, list[Any]]:
        executor_action = self.state.executor_action
        result = self.verifier.run(
            instruction=self.instruction,
            plan=self.state.plan,
            screenshot_b64=screenshot_b64,
            executor_action=executor_action,
        )

        self._task_log.info(
            "\n=== VERIFICATION | decision=%s ===\nReasoning: %s\nDiagnosis: %s\nSuggestion: %s",
            result.decision.value,
            result.reasoning or "(none)",
            result.diagnosis or "(none)",
            result.suggestion or "(none)",
        )

        if result.decision == VerifierDecision.DONE:
            if executor_action and executor_action.decision == ExecutorDecision.FAIL:
                self.state.phase = Phase.FAILED
                return f"Verifier confirms: task infeasible. {result.reasoning}", [{"action_type": "FAIL"}]
            self.state.phase = Phase.DONE
            return f"Verifier: task complete. {result.reasoning}", [{"action_type": "DONE"}]

        if result.decision == VerifierDecision.INFEASIBLE:
            self.state.phase = Phase.FAILED
            return f"Verifier: task infeasible. {result.reasoning}", [{"action_type": "FAIL"}]

        # FAILED — replan or submit best-effort
        if self.state.plan.revision_count >= config.max_replan_count:
            self._task_log.info("\n=== REPLAN BUDGET EXHAUSTED — submitting as DONE ===")
            self.state.phase = Phase.DONE
            return "Replan budget exhausted — submitting.", [{"action_type": "DONE"}]

        self._task_log.info(
            "\n=== REPLAN TRIGGERED ===\nDiagnosis: %s\nSuggestion: %s",
            result.diagnosis or "(none)",
            result.suggestion or "(none)",
        )
        self.state.verification_result = result
        self.state.phase = Phase.REPLANNING
        return "Verification failed — replanning.", [{"action_type": "WAIT"}]

    def _handle_replanning(self, screenshot_b64: str) -> tuple[str, list[Any]]:
        result = self.state.verification_result

        self.planner.reset()
        new_plan = self.planner.run(
            instruction=self.instruction,
            screenshot_b64=screenshot_b64,
            replan=True,
            current_plan=self.state.plan,
            verifier_diagnosis=result.diagnosis if result else "",
            verifier_suggestion=result.suggestion if result else "",
        )
        self.state.plan = new_plan
        self.executor.reset()
        self.state.replan_guidance = result.suggestion if result else ""
        self.state.verification_result = None
        self.state.executor_action = None

        self._task_log.info(
            "\n=== REPLAN (revision %d) ===\n%s",
            new_plan.revision_count,
            new_plan.format(),
        )

        self.state.phase = Phase.EXECUTING
        return "Replanning complete — resuming execution.", [{"action_type": "WAIT"}]

    def log_usage(self) -> None:
        self.gate.log_usage()
        self.planner.log_usage()
        self.executor.log_usage()
        self.verifier.log_usage()
