from __future__ import annotations

import asyncio
import json
import traceback
from dataclasses import dataclass
from typing import Any, Awaitable, Callable, Optional

from .agent_stream_state import ToolRecord, UsageData
from .context import Context
from .model import (
    Message,
    ToolMessageContent,
    FunctionInfo,
    ToolCallInfo,
    ROLE_ASSISTANT,
    ROLE_TOOL,
)
from .tool_runtime_utils import (
    JsonObject,
    parse_tool_arguments,
    parse_tool_arguments_with_status,
    summarize_tool_arguments,
)
from .tools.base import BaseToolRegistry, ToolResult
from .events import (
    EventPayload,
    ev_error,
    ev_assistant_final,
    ev_tool_output,
    ev_tool_exit,
    ev_tool_call_request,
    ev_tool_call_debug_stage,
)


@dataclass(slots=True)
class ToolRuntime:
    tool_registry: BaseToolRegistry
    publish: Callable[[EventPayload], Awaitable[None]]
    execute_remote_tool: Callable[[Context, str, JsonObject], Awaitable[ToolResult]]
    cancel_event: asyncio.Event
    debug_tools: bool = False

    async def _emit_tool_denial(
        self,
        ctx: Context,
        name: str,
        tcid: str,
        message: str | None,
        reasons: list[str] | None,
        args: Optional[JsonObject] = None,
        full_metadata: dict | None = None,
    ) -> None:
        detail = message or f"{name} execution denied."
        info: dict[str, Any] = {}
        try:
            tool_class = self.tool_registry.get_tool_class(name)
            if tool_class and hasattr(tool_class, "on_denial"):
                info = tool_class.on_denial(args or {})
        except Exception:
            pass

        payload_metadata = full_metadata if full_metadata else {"tool": name, "reasons": reasons}
        payload = ToolMessageContent(
            status="failure",
            data=full_metadata,
            error=detail,
            metadata=payload_metadata,
            info=info,
        )
        await self.publish(ev_tool_output(stream="stdout", text=json.dumps(payload.model_dump(), ensure_ascii=False) + "\n"))
        await self.publish(ev_tool_output(stream="stderr", text=detail + "\n"))
        await self.publish(ev_tool_exit(1))
        try:
            ctx.append_message(Message(role=ROLE_TOOL, content=payload, tool_call_id=tcid, name=name))
        except Exception:
            pass

    async def process_tool_records(
        self,
        *,
        ctx: Context,
        tool_records: list[ToolRecord],
        final: str,
        reasoning: str,
        reasoning_details: Any,
        usage_data: UsageData | None,
        model_name: str | None,
    ) -> bool:
        try:
            tool_calls_objects: list[ToolCallInfo] = []
            for i, rec in enumerate(tool_records):
                parsed_args = parse_tool_arguments(rec.get("function", {}).get("arguments", ""))
                function = FunctionInfo(
                    name=rec.get("function", {}).get("name", ""),
                    arguments=parsed_args if parsed_args else None,
                )
                tool_call = ToolCallInfo(
                    id=rec.get("id") or f"call_{i}",
                    type="function",
                    function=function,
                )
                tool_calls_objects.append(tool_call)

            reasoning_content = reasoning if reasoning else None
            ctx.append_message(
                Message(
                    role=ROLE_ASSISTANT,
                    content=final,
                    tool_calls=tool_calls_objects,
                    usage=usage_data,
                    reasoning_content=reasoning_content,
                    model_name=model_name,
                    extra={"reasoning_details": reasoning_details} if reasoning_details else None,
                )
            )
        except Exception as e:
            error_msg = f"Tool record assembly failed: {str(e)}"
            await self.publish(ev_error(error_msg, exception_type=type(e).__name__, traceback_text=traceback.format_exc()))
            await self.publish(ev_assistant_final(f"Error: {error_msg}"))

        for i, rec in enumerate(tool_records):
            name = rec.get("function", {}).get("name", "")
            arg_str = rec.get("function", {}).get("arguments", "")
            if self.debug_tools:
                await self.publish(
                    ev_tool_call_debug_stage(
                        stage="assembled_final",
                        tool=name,
                        id=rec.get("id") or "",
                        index=int(rec.get("index") or 0),
                        args_summary=summarize_tool_arguments(arg_str),
                    )
                )
            rec_index = rec.get("index")
            rec_tcid = rec.get("id")
            if rec_index is not None:
                delta_id = f"idx:{rec_index}"
            elif rec_tcid:
                delta_id = f"id:{rec_tcid}"
            else:
                delta_id = f"call_{i}"
            loop = asyncio.get_running_loop()
            preview_fut: asyncio.Future = loop.create_future()
            await self.publish(ev_tool_call_request(tool=name, delta_id=delta_id, future=preview_fut))
            await preview_fut

        request_new_input = False
        for i, rec in enumerate(tool_records):
            name = rec.get("function", {}).get("name", "")
            arg_str = rec.get("function", {}).get("arguments", "")
            args, parse_ok = parse_tool_arguments_with_status(arg_str)
            tcid = rec.get("id") or f"call_{i}"

            if arg_str and not parse_ok:
                err = ToolMessageContent(
                    status="failure",
                    error="invalid_tool_arguments_json",
                    data={
                        "message": "Tool arguments are not valid JSON object.",
                        "received": {"name": name, "arguments": arg_str},
                    },
                )
                await self.publish(
                    ev_tool_call_debug_stage(
                        stage="invalid_arguments",
                        tool=name,
                        id=tcid,
                        index=i,
                        args_summary=summarize_tool_arguments(arg_str),
                    )
                )
                try:
                    ctx.append_message(Message(role=ROLE_TOOL, content=err, tool_call_id=tcid, name=name))
                except Exception:
                    traceback.print_exc()
                continue

            tool_class = self.tool_registry.get_tool_class(name)
            if tool_class is None:
                if ctx.get_mcp_tool_definition(name):
                    tool_result = await self.execute_remote_tool(ctx, name, args)
                else:
                    available: list[str] = []
                    try:
                        available = self.tool_registry.list_tools()
                    except Exception:
                        available = []
                    err = ToolMessageContent(
                        status="failure",
                        error="unknown_tool",
                        data={
                            "message": f"Unknown tool: {name or '(empty)'}",
                            "received": {"name": name, "arguments": arg_str},
                            "available_tools": [a for a in available if a],
                            "hint": "Ensure the assistant emits tool_calls with a valid function.name that matches a registered tool.",
                        },
                    )
                    try:
                        ctx.append_message(Message(role=ROLE_TOOL, content=err, tool_call_id=tcid, name=name))
                    except Exception:
                        traceback.print_exc()
                    continue
            else:
                try:
                    tool_result = await tool_class.run(
                        args,
                        cancel_event=self.cancel_event,
                        approval_publish=self.publish,
                        event_publish=self.publish,
                    )
                except Exception as e:
                    traceback.print_exc()
                    tool_result = ToolResult(status="failure", error=str(e))

            if tool_result.status == "success":
                try:
                    tool_content = ToolMessageContent(
                        status=tool_result.status,
                        data=tool_result.content,
                        error=tool_result.error,
                        metadata=tool_result.metadata,
                    )
                    ctx.append_message(Message(role=ROLE_TOOL, content=tool_content, tool_call_id=tcid, name=name))
                except Exception:
                    traceback.print_exc()
                continue

            reasons = None
            if isinstance(tool_result.metadata, dict):
                reasons = tool_result.metadata.get("reasons")
            detail = tool_result.error or "Tool execution failed."
            await self._emit_tool_denial(
                ctx,
                name,
                tcid,
                detail,
                reasons,
                args,
                full_metadata=tool_result.metadata if isinstance(tool_result.metadata, dict) else None,
            )
            request_new_input = bool(tool_result.needs_user_input)
        return request_new_input
