from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Any, Callable

from runtime.planner_config import PlannerConfig
from runtime.tool_specs import ACTION_TOOL_SPECS, ToolSpec

READ_TOOL_SPECS = [
    ToolSpec("take_screenshot", "Capture the current device screen.", {"type": "object", "properties": {"device_id": {"type": "string"}}}),
    ToolSpec("describe_screen", "Inspect the current screen and UI.", {"type": "object", "properties": {"device_id": {"type": "string"}}}),
    ToolSpec("find_text", "Find visible text on the current screen.", {"type": "object", "required": ["query"], "properties": {"query": {"type": "string"}, "device_id": {"type": "string"}}}),
    ToolSpec("get_ui_tree", "Read the current accessibility/UI tree.", {"type": "object", "properties": {"device_id": {"type": "string"}, "include_app_info": {"type": "boolean"}}}),
    ToolSpec("list_devices", "List configured devices.", {"type": "object", "properties": {}}),
    ToolSpec("device_status", "Read one device status.", {"type": "object", "required": ["device_id"], "properties": {"device_id": {"type": "string"}}}),
]


@dataclass(frozen=True)
class ChatResult:
    content: str
    tool_calls: int


class ConversationAgent:
    """Small OpenAI-compatible agent loop backed by Host Agent tools."""

    def __init__(
        self,
        *,
        config: PlannerConfig,
        tools: dict[str, Callable[..., Any]],
        max_rounds: int = 8,
        event_logger: Callable[[dict[str, Any]], None] | None = None,
    ) -> None:
        self.config = config
        self.tools = tools
        self.max_rounds = max_rounds
        self.event_logger = event_logger
        self._client: Any | None = None

    def chat(self, messages: list[dict[str, Any]], *, tools: dict[str, Callable[..., Any]] | None = None) -> ChatResult:
        if self.config.provider == "anthropic":
            return self._chat_anthropic(messages, tools=tools)
        if self.config.provider not in {"openai", "openai_compatible"}:
            raise RuntimeError("unsupported conversation provider")
        client = self._get_client()
        active_tools = tools or self.tools
        history: list[dict[str, Any]] = [
            {
                "role": "system",
                "content": (
                    "You are a mobile device assistant. Reply naturally when no "
                    "device action is needed. Call a tool when the user asks you "
                    "to inspect or operate a device. Never claim an action was "
                    "completed unless the tool result confirms it."
                ),
            },
            *[_normalize_message(message) for message in messages],
        ]
        calls = 0
        for _ in range(self.max_rounds):
            response = client.chat.completions.create(
                model=self.config.resolved_model(),
                messages=history,
                tools=[_openai_tool(spec) for spec in [*ACTION_TOOL_SPECS, *READ_TOOL_SPECS]],
                tool_choice="auto",
                parallel_tool_calls=False,
                timeout=self.config.timeout,
            )
            message = response.choices[0].message
            text = getattr(message, "content", None)
            thinking = getattr(message, "reasoning_content", None) or getattr(message, "reasoning", None)
            tool_calls = getattr(message, "tool_calls", None) or []
            self._log({"type": "llm_response", "content": text, "thinking": thinking,
                       "tool_calls": [{"id": c.id, "name": c.function.name, "arguments": c.function.arguments} for c in tool_calls]})
            if not tool_calls:
                self._log({"type": "final_reply", "content": str(text or "")})
                return ChatResult(content=str(text or ""), tool_calls=calls)
            history.append(_message_dict(message))
            for tool_call in tool_calls:
                name = tool_call.function.name
                arguments: dict[str, Any] = {}
                try:
                    decoded_arguments = json.loads(tool_call.function.arguments or "{}")
                    if not isinstance(decoded_arguments, dict):
                        raise ValueError("tool arguments must be an object")
                    arguments = decoded_arguments
                    arguments.pop("purpose", None)
                    arguments.pop("expected_outcome", None)
                    result = active_tools[name](**arguments)
                except Exception as exc:
                    result = {"ok": False, "error": str(exc)}
                calls += 1
                self._log({"type": "tool_result", "tool_call_id": tool_call.id,
                           "tool_name": name, "arguments": arguments, "result": result})
                history.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "content": json.dumps(result, ensure_ascii=True, default=str),
                    }
                )
        raise RuntimeError("conversation exceeded the maximum tool-call rounds")

    def _chat_anthropic(self, messages: list[dict[str, Any]], *, tools: dict[str, Callable[..., Any]] | None = None) -> ChatResult:
        import anthropic

        kwargs: dict[str, Any] = {}
        if self.config.api_key:
            kwargs["api_key"] = self.config.api_key
        if self.config.base_url:
            kwargs["base_url"] = self.config.base_url
        client = anthropic.Anthropic(**kwargs)
        history = [_normalize_anthropic_message(message) for message in messages]
        active_tools = tools or self.tools
        calls = 0
        for _ in range(self.max_rounds):
            response = client.messages.create(
                model=self.config.resolved_model(),
                max_tokens=2048,
                system="You are a mobile device assistant. Reply naturally, or call one device tool when needed. Never claim an action succeeded unless its tool result confirms it.",
                messages=history,
                tools=[_anthropic_tool(spec) for spec in [*ACTION_TOOL_SPECS, *READ_TOOL_SPECS]],
                timeout=self.config.timeout,
            )
            blocks = getattr(response, "content", []) or []
            text_parts = [getattr(block, "text", "") for block in blocks if getattr(block, "type", None) == "text"]
            thinking_parts = [getattr(block, "thinking", "") for block in blocks if getattr(block, "type", None) == "thinking"]
            uses = [block for block in blocks if getattr(block, "type", None) == "tool_use"]
            self._log({"type": "llm_response", "content": "\n".join(x for x in text_parts if x), "thinking": "\n".join(x for x in thinking_parts if x), "tool_calls": [{"id": u.id, "name": u.name, "arguments": u.input} for u in uses]})
            if not uses:
                content = "\n".join(x for x in text_parts if x)
                self._log({"type": "final_reply", "content": content})
                return ChatResult(content=content, tool_calls=calls)
            history.append({"role": "assistant", "content": [_anthropic_block_dict(block) for block in blocks if getattr(block, "type", None) != "thinking"]})
            results = []
            for use in uses:
                arguments = dict(use.input) if isinstance(use.input, dict) else {}
                try:
                    result = active_tools[use.name](**arguments)
                except Exception as exc:
                    result = {"ok": False, "error": str(exc)}
                calls += 1
                self._log({"type": "tool_result", "tool_call_id": use.id, "tool_name": use.name, "arguments": arguments, "result": result})
                results.append({"type": "tool_result", "tool_use_id": use.id, "content": json.dumps(result, ensure_ascii=True, default=str)})
            history.append({"role": "user", "content": results})
        raise RuntimeError("conversation exceeded the maximum tool-call rounds")

    def chat_for_device(self, device_id: str, messages: list[dict[str, Any]]) -> ChatResult:
        """Run a conversation bound to exactly one device."""
        device_tools = {
            name: _bind_device(tool, device_id, name=name)
            for name, tool in self.tools.items()
        }
        return self.chat(messages, tools=device_tools)

    def _log(self, event: dict[str, Any]) -> None:
        if self.event_logger is not None:
            try:
                self.event_logger(event)
            except Exception:
                pass

    def _get_client(self) -> Any:
        if self._client is None:
            from openai import OpenAI

            kwargs: dict[str, Any] = {}
            if self.config.api_key:
                kwargs["api_key"] = self.config.api_key
            if self.config.base_url:
                kwargs["base_url"] = self.config.base_url
            self._client = OpenAI(**kwargs)
        return self._client


def _openai_tool(spec: ToolSpec) -> dict[str, Any]:
    parameters = dict(spec.parameters)
    properties = dict(parameters.get("properties") or {})
    properties.setdefault(
        "device_id",
        {"type": "string", "description": "Target device ID when needed."},
    )
    parameters["properties"] = properties
    return {
        "type": "function",
        "function": {
            "name": spec.name,
            "description": spec.description,
            "parameters": parameters,
        },
    }


def _anthropic_tool(spec: ToolSpec) -> dict[str, Any]:
    return {"name": spec.name, "description": spec.description, "input_schema": spec.parameters}


def _anthropic_block_dict(block: Any) -> dict[str, Any]:
    block_type = getattr(block, "type", "")
    if block_type == "text":
        return {"type": "text", "text": getattr(block, "text", "")}
    if block_type == "thinking":
        return {"type": "thinking", "thinking": getattr(block, "thinking", "")}
    return {"type": "tool_use", "id": block.id, "name": block.name, "input": block.input}


def _normalize_anthropic_message(message: dict[str, Any]) -> dict[str, Any]:
    normalized = _normalize_message(message)
    content = normalized["content"]
    if isinstance(content, str):
        return normalized
    blocks = []
    for block in content:
        if block.get("type") == "text":
            blocks.append(block)
        elif block.get("type") == "image_url":
            url = block.get("image_url", {}).get("url", "")
            if isinstance(url, str) and url.startswith("data:"):
                header, data = url.split(",", 1)
                media_type = header[5:].split(";", 1)[0]
                blocks.append({"type": "image", "source": {"type": "base64", "media_type": media_type, "data": data}})
    return {"role": normalized["role"], "content": blocks}


def _message_dict(message: Any) -> dict[str, Any]:
    result: dict[str, Any] = {"role": "assistant"}
    content = getattr(message, "content", None)
    if content is not None:
        result["content"] = content
    calls = getattr(message, "tool_calls", None) or []
    result["tool_calls"] = [
        {
            "id": call.id,
            "type": "function",
            "function": {
                "name": call.function.name,
                "arguments": call.function.arguments,
            },
        }
        for call in calls
    ]
    return result


def _normalize_message(message: dict[str, Any]) -> dict[str, Any]:
    """Accept OpenAI text/image content blocks and a compact image_base64 form."""
    role = message.get("role")
    content = message.get("content")
    if isinstance(content, str):
        return {"role": role, "content": content}
    if isinstance(content, list):
        blocks: list[dict[str, Any]] = []
        for block in content:
            if not isinstance(block, dict):
                continue
            if block.get("type") == "text" and isinstance(block.get("text"), str):
                blocks.append({"type": "text", "text": block["text"]})
            elif block.get("type") == "image_url":
                image_url = block.get("image_url")
                if isinstance(image_url, str):
                    image_url = {"url": image_url}
                if isinstance(image_url, dict) and isinstance(image_url.get("url"), str):
                    blocks.append({"type": "image_url", "image_url": {"url": image_url["url"]}})
        if blocks:
            return {"role": role, "content": blocks}
    image = message.get("image_base64")
    if isinstance(image, str) and image:
        mime = str(message.get("mime_type") or "image/png")
        text = message.get("text")
        blocks = []
        if isinstance(text, str) and text:
            blocks.append({"type": "text", "text": text})
        blocks.append({"type": "image_url", "image_url": {"url": f"data:{mime};base64,{image}"}})
        return {"role": role, "content": blocks}
    raise ValueError("message content must be text, image blocks, or image_base64")


def _bind_device(tool: Callable[..., Any], device_id: str, *, name: str) -> Callable[..., Any]:
    def bound(**arguments: Any) -> Any:
        if name == "list_devices":
            return tool(**arguments)
        requested = arguments.get("device_id")
        if requested is not None and requested != device_id:
            raise ValueError(f"conversation is bound to device {device_id}")
        arguments["device_id"] = device_id
        return tool(**arguments)

    return bound