"""Feasibility gate — upfront pre-check that determines whether a task is feasible before execution."""

from __future__ import annotations

import time
from typing import Any

from .config import config
from .knowledge import KnowledgeBank
from .llm_client import LLMClient
from .prompts import FEASIBILITY_GATE_SYSTEM_PROMPT
from .task_logger import get_agent_logger, log_agent_api_call, save_screenshot_bytes
from .tools import (
    CLICK,
    HOLD_AND_PRESS,
    HOTKEY,
    LIST_WINDOWS,
    OPEN,
    PROBE_BASH,
    PROBE_CHROME_CDP,
    PROBE_PYTHON,
    READ_FILE,
    REPORT_FEASIBLE,
    REPORT_INFEASIBLE,
    SCREENSHOT,
    SCROLL,
    SWITCH_APPLICATIONS,
    TYPE,
    WAIT,
    WEB_FETCH,
    WEB_SEARCH,
    read_file_on_vm,
    run_bash_on_vm,
    run_cdp_on_vm,
    run_python_on_vm,
    take_screenshot_on_vm,
    web_extract,
    web_search,
)
from .types import FeasibilityResult
from .utils import ToolResult

_task_log = get_agent_logger("feasibility_gate")

GATE_TOOLS = [
    PROBE_BASH,
    PROBE_PYTHON,
    PROBE_CHROME_CDP,
    READ_FILE,
    LIST_WINDOWS,
    SCREENSHOT,
    WEB_SEARCH,
    WEB_FETCH,
    CLICK,
    TYPE,
    SCROLL,
    HOTKEY,
    HOLD_AND_PRESS,
    WAIT,
    OPEN,
    SWITCH_APPLICATIONS,
    REPORT_FEASIBLE,
    REPORT_INFEASIBLE,
]

_GUI_TOOL_RESULT_TEXT = "GUI action executed. A fresh screenshot follows."

# MARK: Key map (shared with executor)

_KEY_MAP = {
    "page_down": "pagedown",
    "page_up": "pageup",
    "super_l": "win",
    "super": "command",
    "escape": "esc",
}

_SWITCH_APPLICATIONS_TEMPLATE = """import subprocess
import difflib
import pyautogui
import time
pyautogui.press('escape')
time.sleep(0.5)
output = subprocess.check_output(['wmctrl', '-lx']).decode('utf-8').splitlines()
window_titles = [line.split(None, 4)[2] for line in output]
closest_matches = difflib.get_close_matches({app_code!r}, window_titles, n=1, cutoff=0.1)
if closest_matches:
    closest_match = closest_matches[0]
    window_id = None
    for line in output:
        if closest_match in line:
            window_id = line.split()[0]
            break
    if window_id:
        subprocess.run(['wmctrl', '-ia', window_id])
        subprocess.run(['wmctrl', '-ir', window_id, '-b', 'add,maximized_vert,maximized_horz'])
        time.sleep(5)
"""


# MARK: Tool handlers


def _exec_gui_on_vm(gate: FeasibilityGate, pycode: str) -> ToolResult:
    """Execute pyautogui code on the VM and return a screenshot as the result."""
    run_python_on_vm(gate.env, pycode)
    time.sleep(1)
    result, png_bytes = take_screenshot_on_vm(gate.env)
    if png_bytes:
        save_screenshot_bytes(png_bytes, label="gate_gui")
    return result if not result.error else ToolResult(output=_GUI_TOOL_RESULT_TEXT)


def _handle_probe_bash(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    result = run_bash_on_vm(gate.env, inp.get("script", ""), timeout=inp.get("timeout", 30))
    return None, result


def _handle_probe_python(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    result = run_python_on_vm(gate.env, inp.get("code", ""))
    return None, result


def _handle_probe_chrome(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    result = run_cdp_on_vm(gate.env, inp.get("code", ""))
    return None, result


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


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


def _handle_screenshot(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    result, png_bytes = take_screenshot_on_vm(gate.env)
    if png_bytes:
        save_screenshot_bytes(png_bytes, label="gate_screenshot")
    return None, result


def _handle_web_search(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    result = web_search(inp.get("query", ""), inp.get("max_results", 5))
    return None, result


def _handle_web_fetch(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    result = web_extract(inp.get("urls", []), inp.get("max_chars", 8000))
    return None, result


def _handle_report_feasible(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    reason = inp.get("reason", "")
    _task_log.info("Gate verdict: FEASIBLE — %s", reason)
    return FeasibilityResult(infeasible=False, reason=reason), ToolResult(output="Verdict recorded: feasible.")


def _handle_report_infeasible(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    reason = inp.get("reason", "")
    _task_log.info("Gate verdict: INFEASIBLE — %s", reason)
    return FeasibilityResult(infeasible=True, reason=reason), ToolResult(output="Verdict recorded: infeasible.")


def _scale_xy(coordinate: list[int], resize_factor: tuple[float, float]) -> tuple[int, int]:
    return int(coordinate[0] * resize_factor[0]), int(coordinate[1] * resize_factor[1])


def _handle_click(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    coordinate = inp.get("coordinate") or [0, 0]
    num_clicks = inp.get("num_clicks", 1)
    button_type = inp.get("button_type", "left")
    hold_keys = inp.get("hold_keys", []) or []
    if isinstance(hold_keys, str):
        hold_keys = [hold_keys]
    x, y = _scale_xy(coordinate, gate.resize_factor)
    pycode = "import pyautogui; "
    for key in hold_keys:
        k = _KEY_MAP.get(str(key).lower(), str(key).lower())
        pycode += f"pyautogui.keyDown({k!r}); "
    pycode += f"pyautogui.click(x={x}, y={y}, clicks={num_clicks}, button={button_type!r}); "
    for key in hold_keys:
        k = _KEY_MAP.get(str(key).lower(), str(key).lower())
        pycode += f"pyautogui.keyUp({k!r}); "
    return None, _exec_gui_on_vm(gate, pycode)


def _handle_type(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    text = inp.get("text", "") or ""
    coordinate = inp.get("coordinate")
    overwrite = bool(inp.get("overwrite", False))
    enter = bool(inp.get("enter", False))
    pycode = "import pyautogui; import time; "
    if coordinate:
        x, y = _scale_xy(coordinate, gate.resize_factor)
        pycode += f"pyautogui.click(x={x}, y={y}); time.sleep(0.5); "
    if overwrite:
        pycode += "pyautogui.hotkey('ctrl', 'a'); pyautogui.press('backspace'); "
    pycode += f"pyautogui.write({text!r}, interval=0.05); "
    if enter:
        pycode += "pyautogui.press('enter'); "
    return None, _exec_gui_on_vm(gate, pycode)


def _handle_scroll(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    coordinate = inp.get("coordinate") or [0, 0]
    clicks = int(inp.get("clicks", 0))
    shift = bool(inp.get("shift", False))
    x, y = _scale_xy(coordinate, gate.resize_factor)
    pycode = "import pyautogui; import time; "
    pycode += f"pyautogui.moveTo(x={x}, y={y}); time.sleep(0.3); "
    if shift:
        pycode += f"pyautogui.hscroll({clicks}); "
    else:
        pycode += f"pyautogui.vscroll({clicks}); "
    return None, _exec_gui_on_vm(gate, pycode)


def _handle_hotkey(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    keys = inp.get("keys", []) or []
    if isinstance(keys, str):
        keys = [keys]
    normalized = [_KEY_MAP.get(str(k).lower(), str(k).lower()) for k in keys]
    keys_repr = ", ".join(repr(k) for k in normalized)
    pycode = f"import pyautogui; pyautogui.hotkey({keys_repr}); "
    return None, _exec_gui_on_vm(gate, pycode)


def _handle_hold_and_press(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    hold_keys = inp.get("hold_keys", []) or []
    press_keys = inp.get("press_keys", []) or []
    if isinstance(hold_keys, str):
        hold_keys = [hold_keys]
    if isinstance(press_keys, str):
        press_keys = [press_keys]
    holds = [_KEY_MAP.get(str(k).lower(), str(k).lower()) for k in hold_keys]
    presses = [_KEY_MAP.get(str(k).lower(), str(k).lower()) for k in press_keys]
    pycode = "import pyautogui; "
    for k in holds:
        pycode += f"pyautogui.keyDown({k!r}); "
    for k in presses:
        pycode += f"pyautogui.press({k!r}); "
    for k in holds:
        pycode += f"pyautogui.keyUp({k!r}); "
    return None, _exec_gui_on_vm(gate, pycode)


def _handle_wait(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    seconds = float(inp.get("seconds", 1.0))
    time.sleep(seconds)
    return None, ToolResult(output=f"Waited {seconds}s.")


def _handle_open(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    name = str(inp.get("app_or_filename", ""))
    pycode = "import pyautogui; import time; "
    pycode += "pyautogui.hotkey('win'); time.sleep(1); "
    pycode += f"pyautogui.write({name!r}); time.sleep(1); "
    pycode += "pyautogui.press('enter'); time.sleep(5); "
    return None, _exec_gui_on_vm(gate, pycode)


def _handle_switch_applications(gate: FeasibilityGate, inp: dict) -> tuple[FeasibilityResult | None, ToolResult]:
    app_code = str(inp.get("app_code", ""))
    pycode = _SWITCH_APPLICATIONS_TEMPLATE.format(app_code=app_code)
    run_python_on_vm(gate.env, pycode)
    time.sleep(1)
    result, png_bytes = take_screenshot_on_vm(gate.env)
    if png_bytes:
        save_screenshot_bytes(png_bytes, label="gate_switch")
    return None, result if not result.error else ToolResult(output=_GUI_TOOL_RESULT_TEXT)


_TOOL_DISPATCH: dict[str, Any] = {
    "probe_bash": _handle_probe_bash,
    "probe_python": _handle_probe_python,
    "probe_chrome_cdp": _handle_probe_chrome,
    "read_file": _handle_read_file,
    "list_windows": _handle_list_windows,
    "screenshot": _handle_screenshot,
    "web_search": _handle_web_search,
    "web_fetch": _handle_web_fetch,
    "click": _handle_click,
    "type": _handle_type,
    "scroll": _handle_scroll,
    "hotkey": _handle_hotkey,
    "hold_and_press": _handle_hold_and_press,
    "wait": _handle_wait,
    "open": _handle_open,
    "switch_applications": _handle_switch_applications,
    "report_feasible": _handle_report_feasible,
    "report_infeasible": _handle_report_infeasible,
}


# MARK: FeasibilityGate


class FeasibilityGate:
    """Pre-execution gate that probes the VM to determine task feasibility.

    Runs as a single multi-turn LLM call (like the Verifier). If it concludes
    the task is infeasible, the orchestrator skips execution entirely.
    """

    def __init__(
        self,
        env: Any,
        knowledge: KnowledgeBank | None = None,
        screen_size: tuple[int, int] = (1920, 1080),
    ) -> None:
        self.env = env
        self.knowledge = knowledge
        self.resize_factor = (screen_size[0] / 1280, screen_size[1] / 720)

        self.llm = LLMClient(
            name="feasibility_gate",
            system_prompt=FEASIBILITY_GATE_SYSTEM_PROMPT,
            tools=[t.schema for t in GATE_TOOLS],
            model=config.gate_model,
            provider=config.provider,
            only_n_most_recent_images=config.gate_n_recent_images,
            effort=config.gate_effort,
        )

        global _task_log
        _task_log = get_agent_logger("feasibility_gate")

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

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

    def run(self, instruction: str, screenshot_b64: str) -> FeasibilityResult:
        """Probe the VM and return a feasibility verdict."""
        self.llm.reset()

        prompt = self._build_initial_prompt(instruction)
        images = [screenshot_b64] if screenshot_b64 else None
        if images:
            prompt += "\nThe current desktop screenshot is attached."
        self.llm.add_user_message(text=prompt, images=images)

        _task_log.info("\n=== FEASIBILITY GATE START ===\nTask: %s", instruction)

        for step in range(config.gate_max_steps):
            verdict = self._execute_step()
            if verdict is not None:
                _task_log.info(
                    "\n=== FEASIBILITY GATE END | infeasible=%s ===\nReason: %s",
                    verdict.infeasible,
                    verdict.reason,
                )
                return verdict

        _task_log.info(
            "\nBudget exhausted after %d iterations — defaulting to feasible.",
            config.gate_max_steps,
        )
        return FeasibilityResult(infeasible=False, reason="Budget exhausted — defaulting to feasible.")

    def _execute_step(self) -> FeasibilityResult | None:
        """Dispatch a single LLM turn's tool calls. Returns verdict if a report tool 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: FeasibilityResult | 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="feasibility_gate",
            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

    def _build_initial_prompt(self, instruction: str) -> str:
        parts = [f"TASK: {instruction}"]

        if self.knowledge:
            kb = self.knowledge.format_for_prompt()
            if kb:
                parts.append(f"\n{kb}")

        parts.append(
            "\nProbe the VM with the available tools, then call report_feasible or report_infeasible exactly once."
        )
        return "\n".join(parts)
