from __future__ import annotations

from typing import TYPE_CHECKING, Any

from core.errors import TaskFailedError
from core.models import Scene
from runtime.context import TaskContext
from runtime.planner import PlannedStep, Planner
from runtime.planner_config import PlannerConfig, load_config
from runtime.planner_prompts import PLANNER_SYSTEM_PROMPT, planner_user_prompt
from runtime.tool_calling_client import ToolCallingClient, build_client
from runtime.tool_specs import ALL_TOOL_SPECS

if TYPE_CHECKING:
    from world.models import WorldState

FINISH_TASK_TOOL = "finish_task"


class AIPlanner(Planner):
    def __init__(
        self,
        *,
        client: ToolCallingClient | None = None,
        config: PlannerConfig | None = None,
    ) -> None:
        self.config = config or load_config()
        self.client = client or build_client(self.config)

    def plan(
        self,
        *,
        goal: str,
        scene: Scene,
        context: TaskContext,
        world: "WorldState | None" = None,
        screenshot: bytes | None = None,
    ) -> list[PlannedStep]:
        scene_json = _without_ocr(scene.to_dict()) if self.config.multimodal else scene.to_dict()
        user_prompt = planner_user_prompt(
            goal=goal,
            scene_json=scene_json,
            history_summary=_history_summary(world),
            device_platform=context.device_platform,
        )
        decision = self.client.decide(
            system_prompt=PLANNER_SYSTEM_PROMPT,
            user_prompt=user_prompt,
            screenshot=screenshot,
            tools=ALL_TOOL_SPECS,
            timeout=self.config.timeout,
        )

        if decision.tool_name == FINISH_TASK_TOOL:
            if decision.arguments.get("success"):
                return []
            raise TaskFailedError(decision.arguments.get("reason") or "task failed")

        return [
            PlannedStep(
                action=decision.tool_name,
                description=(
                    f"AI planner: {decision.purpose}"
                    if decision.purpose
                    else f"AI planner: {decision.tool_name}({decision.arguments})"
                ),
                args=dict(decision.arguments),
                expected_text=decision.expected_outcome,
                purpose=decision.purpose,
                expected_outcome=decision.expected_outcome,
                prompt=decision.user_prompt or user_prompt,
                rationale=decision.text_output,
                thinking=decision.thinking,
            )
        ]

    def goal_reached(self, *, goal: str, scene: Scene, context: TaskContext) -> bool:
        # Completion is signaled exclusively via the finish_task tool call
        # (mapped to an empty plan above), never via this hook.
        return False


def _history_summary(world: "WorldState | None") -> list[dict[str, Any]]:
    if world is None:
        return []
    return [
        {
            "page": event.page,
            "action": event.action,
            "arguments": dict(event.arguments),
            "rationale": event.rationale,
            "purpose": event.purpose,
            "expected_outcome": event.expected_outcome,
            "success": event.success,
        }
        for event in world.history
    ]


def _without_ocr(scene_json: dict[str, Any]) -> dict[str, Any]:
    cleaned = dict(scene_json)
    elements = cleaned.get("elements")
    if isinstance(elements, list):
        cleaned["elements"] = [
            {key: value for key, value in element.items() if key not in {"source", "confidence", "foreground_color", "background_color"}}
            for element in elements
            if isinstance(element, dict) and element.get("source") != "ocr"
        ]
    cleaned.pop("ocr_elements", None)
    return cleaned