"""Verifier agent — goal-oriented end-state verification via VM probing."""

from __future__ import annotations

from typing import Any

from .config import config
from .llm_client import LLMClient
from .prompts import VERIFIER_SYSTEM_PROMPT
from .task_logger import get_agent_logger, log_agent_api_call, save_screenshot_bytes
from .tools import (
    LIST_WINDOWS,
    READ_FILE,
    RUN_BASH,
    RUN_PYTHON,
    SCREENSHOT,
    SUBMIT_VERDICT,
    read_file_on_vm,
    run_bash_on_vm,
    run_python_on_vm,
    take_screenshot_on_vm,
)
from .types import (
    ExecutorAction,
    ExecutorDecision,
    Plan,
    VerificationResult,
    VerifierDecision,
)
from .utils import ToolResult

_task_log = get_agent_logger("verifier")
VERIFIER_TOOLS = [
    LIST_WINDOWS,
    READ_FILE,
    RUN_BASH,
    RUN_PYTHON,
    SCREENSHOT,
    SUBMIT_VERDICT,
]


# MARK: Tool Handlers


def _handle_run_bash(verifier: Verifier, inp: dict) -> tuple[VerificationResult | None, ToolResult]:
    result = run_bash_on_vm(verifier.env, inp.get("script", ""), timeout=inp.get("timeout", 30))
    return None, result


def _handle_run_python(verifier: Verifier, inp: dict) -> tuple[VerificationResult | None, ToolResult]:
    result = run_python_on_vm(verifier.env, inp.get("code", ""))
    return None, result


def _handle_read_file(verifier: Verifier, inp: dict) -> tuple[VerificationResult | None, ToolResult]:
    result = read_file_on_vm(verifier.env, inp.get("path", ""))
    return None, result


def _handle_list_windows(verifier: Verifier, inp: dict) -> tuple[VerificationResult | None, ToolResult]:
    result = run_bash_on_vm(verifier.env, "wmctrl -lx", timeout=10)
    return None, result


def _handle_screenshot(verifier: Verifier, inp: dict) -> tuple[VerificationResult | None, ToolResult]:
    result, png_bytes = take_screenshot_on_vm(verifier.env)
    saved_path = save_screenshot_bytes(png_bytes, label="verifier_screenshot") if png_bytes else ""
    if result.error:
        _task_log.info("Tool ERROR [screenshot]: %s", result.error)
    else:
        _task_log.info("Tool result [screenshot]: captured %s", saved_path)
    return None, result


def _handle_submit_verdict(verifier: Verifier, inp: dict) -> tuple[VerificationResult | None, ToolResult]:
    verdict = _parse_verdict(inp)
    if verdict is None:
        err = (
            "Invalid verdict. decision must be 'done' | 'failed' | "
            "'infeasible'; when decision='failed', both 'diagnosis' "
            "and 'suggestion' must be non-empty. Please re-submit."
        )
        return None, ToolResult(error=err)

    output = (
        f"Tool result [submit_verdict]: decision={verdict.decision.value}\n"
        f"  reasoning: {verdict.reasoning or '(none)'}\n"
        f"  diagnosis: {verdict.diagnosis or '(none)'}\n"
        f"  suggestion: {verdict.suggestion or '(none)'}"
    )
    return verdict, ToolResult(output=output)


_TOOL_DISPATCH: dict[str, Any] = {
    "run_bash": _handle_run_bash,
    "run_python": _handle_run_python,
    "read_file": _handle_read_file,
    "list_windows": _handle_list_windows,
    "screenshot": _handle_screenshot,
    "submit_verdict": _handle_submit_verdict,
}


# MARK: Verifier


class Verifier:
    """End-state verifier. Probes the VM and returns DONE / FAILED / INFEASIBLE.

    Unlike the prior action-review verifier, this one never inspects the
    executor's trajectory. It reads files, runs shell commands, and queries
    application state — the same way the benchmark grader does — then emits
    a single verdict.
    """

    def __init__(
        self,
        env: Any,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        self.env = env
        self.llm = LLMClient(
            name="verifier",
            system_prompt=VERIFIER_SYSTEM_PROMPT,
            tools=[t.schema for t in VERIFIER_TOOLS],
            model=config.verifier_model,
            provider=config.provider,
            only_n_most_recent_images=config.verifier_n_recent_images,
            effort=config.verifier_effort,
        )

        global _task_log
        _task_log = get_agent_logger("verifier")

    def reset(self) -> None:
        self.llm.reset()

    def log_usage(self) -> None:
        _task_log.info(self.llm.print_usage())

    def run(
        self,
        instruction: str,
        plan: Plan,
        screenshot_b64: str,
        executor_action: ExecutorAction | None = None,
    ) -> VerificationResult:
        """Probe the VM until a verdict is reached."""
        self.llm.reset()

        prompt = self._build_initial_prompt(instruction, plan, executor_action)
        images = None
        if screenshot_b64:
            images = [screenshot_b64]
            prompt += "\nThe current screenshot is attached for context only."
        self.llm.add_user_message(text=prompt, images=images)

        _task_log.info(
            "\n=== VERIFIER START ===\nTask: %s\nPlan:\n%s",
            instruction,
            plan.format(),
        )

        budget = config.verifier_max_steps
        for _ in range(budget):
            verdict = self._execute_step()
            if verdict is not None:
                _task_log.info("\n=== VERIFIER END | decision=%s ===", verdict.decision.value)
                return verdict

        _task_log.info(
            "\nBudget exhausted after %d iterations — entering forced verdict phase.",
            config.verifier_max_steps,
        )
        return self._handle_budget_exhausted()

    def _execute_step(self) -> VerificationResult | None:
        """Dispatch a single LLM turn's tool calls. Returns verdict if submit_verdict fired."""
        tool_calls, text_blocks, thinking_blocks, _ = self.llm.generate_and_process()

        if not tool_calls:
            return None

        tool_call_results: list[tuple[str, ToolResult]] = []
        final_verdict: VerificationResult | None = None

        for i, tc in enumerate(tool_calls):
            handler = _TOOL_DISPATCH.get(tc["name"])
            if handler is None:
                tool_call_results.append((tc["id"], ToolResult(error=f"Unknown tool: {tc['name']}")))
                continue

            verdict, result = handler(self, tc["input"])
            tool_call_results.append((tc["id"], result))

            if verdict is not None:
                final_verdict = verdict
                for remaining in tool_calls[i + 1 :]:
                    tool_call_results.append(
                        (
                            remaining["id"],
                            ToolResult(error="Not executed: a prior tool returned control."),
                        )
                    )
                break

        log_agent_api_call(
            agent_name="verifier",
            tool_calls=tool_calls,
            text_blocks=text_blocks,
            thinking_blocks=thinking_blocks,
            tool_call_results=tool_call_results,
            agent_action=final_verdict,
        )

        self.llm.add_tool_results(tool_call_results)
        return final_verdict

    @staticmethod
    def _build_initial_prompt(instruction: str, plan: Plan, executor_action: ExecutorAction | None = None) -> str:
        parts = [
            f"TASK: {instruction}",
            ("\nPLAN (executor's milestones, for context only — do NOT grade them):\n" + plan.format()),
        ]

        if executor_action and executor_action.decision == ExecutorDecision.FAIL:
            parts.append(
                f"\nEXECUTOR CLAIM: The executor believes this task is INFEASIBLE."
                f"\nExecutor's reason: {executor_action.reason}"
                f"\n\nYour job: Verify whether the task is genuinely infeasible for "
                f"the stated reason. Probe the VM to confirm the missing "
                f"capability/resource. If the task IS feasible, verdict is 'failed' "
                f"with a concrete suggestion for how to accomplish it."
            )
        elif executor_action and executor_action.decision == ExecutorDecision.DONE:
            parts.append(
                f"\nEXECUTOR CLAIM: The executor believes this task is COMPLETE."
                f"\nExecutor's summary: {executor_action.reason}"
                f"\n\nYour job: (1) Verify the task was actually completed — probe "
                f"the VM for evidence. (2) Check whether this looks like a workaround "
                f"for an infeasible task — if the executor produced a substitute, "
                f"stub, or imitation instead of the real thing, the verdict is "
                f"'infeasible'."
            )

        parts.append(
            "\nProbe the VM with run_bash / run_python / read_file / GUI actions "
            "(screenshot), then call submit_verdict once with your decision."
        )
        return "\n".join(parts)


def _parse_verdict(inp: dict[str, Any]) -> VerificationResult | None:
    decision_str = inp.get("decision", "")
    try:
        decision = VerifierDecision(decision_str)
    except ValueError:
        return None

    diagnosis = (inp.get("diagnosis") or "").strip()
    suggestion = (inp.get("suggestion") or "").strip()

    if decision == VerifierDecision.FAILED and (not diagnosis or not suggestion):
        return None

    return VerificationResult(
        decision=decision,
        reasoning=(inp.get("reasoning") or "").strip(),
        diagnosis=diagnosis,
        suggestion=suggestion,
    )
