from __future__ import annotations

import json
from typing import Any, Callable, Dict, List, Optional, Sequence, Union

from .errors import (
    RunnerMaxIterationsError,
    ToolArgumentsError,
    ToolExecutionError,
    ToolNotFoundError,
)
from .registry import ToolRegistry
from .types import InputItem, InputList, RunnerResult

EventCallback = Callable[[str, Dict[str, Any]], None]


def _is_function_call_item(item: Any) -> bool:
    return getattr(item, "type", None) == "function_call" or (
        isinstance(item, dict) and item.get("type") == "function_call"
    )


def _get_item_field(item: Any, key: str) -> Any:
    if isinstance(item, dict):
        return item.get(key)
    return getattr(item, key)


class ToolRunner:
    """Executes the tool-calling loop for OpenAI Responses API."""

    def __init__(
        self,
        client: Any,
        *,
        model: str,
        registry: ToolRegistry,
        max_iters: int = 8,
        max_retry: int = 0,
        parallel_tool_calls: Optional[bool] = None,
        on_event: Optional[EventCallback] = None,
    ) -> None:
        self._client = client
        self._model = model
        self._registry = registry
        self._max_iters = max_iters
        self._max_retry = max_retry
        self._parallel_tool_calls = parallel_tool_calls
        self._on_event = on_event

    def run(
        self,
        prompt_or_input: Union[str, Sequence[InputItem]],
        *,
        tool_choice: Optional[Any] = None,
        instructions: Optional[str] = None,
        extra: Optional[Dict[str, Any]] = None,
    ) -> RunnerResult:
        """Run until the model stops requesting function calls."""
        input_items: InputList = self._normalize_input(prompt_or_input)
        tools_payload = self._registry.as_openai_tools()

        request_kwargs: Dict[str, Any] = {
            "model": self._model,
            "input": list(input_items),
            "tools": tools_payload,
        }
        if tool_choice is not None:
            request_kwargs["tool_choice"] = tool_choice
        if instructions:
            request_kwargs["instructions"] = instructions
        if self._parallel_tool_calls is not None:
            request_kwargs["parallel_tool_calls"] = self._parallel_tool_calls
        if extra:
            request_kwargs.update(extra)

        last_response: Any = None
        consecutive_errors = 0

        for iteration in range(self._max_iters):
            self._emit("llm.request", {"iteration": iteration})
            last_response = self._client.responses.create(**request_kwargs)

            # Important: append model output items to the running input, so the
            # model can "see" what it decided (tool calls, reasoning items, etc).
            output_items = list(getattr(last_response, "output", []) or [])
            input_items.extend(_to_dict_items(output_items))

            tool_calls = [it for it in output_items if _is_function_call_item(it)]
            self._emit(
                "llm.response", {"iteration": iteration, "tool_calls": len(tool_calls)}
            )

            if not tool_calls:
                text = getattr(last_response, "output_text", "") or ""
                return RunnerResult(
                    text=text,
                    input_items=tuple(input_items),
                    last_response=last_response,
                )

            has_success = False
            for call_item in tool_calls:
                call_id = _get_item_field(call_item, "call_id")
                tool_name = _get_item_field(call_item, "name")
                arguments_json = _get_item_field(call_item, "arguments") or "{}"

                self._emit(
                    "tool.call",
                    {
                        "call_id": call_id,
                        "name": tool_name,
                        "arguments": arguments_json,
                    },
                )

                output_str = ""
                try:
                    try:
                        raw_args = json.loads(arguments_json)
                        if not isinstance(raw_args, dict):
                            raise ToolArgumentsError(
                                "Tool arguments must be a JSON object"
                            )
                    except json.JSONDecodeError as exc:
                        raise ToolArgumentsError(
                            f"Invalid JSON arguments for tool '{tool_name}': {exc}"
                        ) from exc

                    tool_def = self._registry.get(tool_name)

                    try:
                        output_str = tool_def.call_and_serialize(raw_args)
                        has_success = True
                    except (ToolArgumentsError, ToolExecutionError):
                        raise
                    except Exception as exc:
                        # Safety net: normalize unknown exceptions.
                        raise ToolExecutionError(
                            tool_name, str(exc), cause=exc
                        ) from exc

                except (
                    ToolArgumentsError,
                    ToolExecutionError,
                    ToolNotFoundError,
                ) as exc:
                    if consecutive_errors < self._max_retry:
                        output_str = f"Error: {exc}"
                    else:
                        raise

                # The model expects a function_call_output item with the same call_id.
                input_items.append(
                    {
                        "type": "function_call_output",
                        "call_id": call_id,
                        "output": output_str,
                    }
                )
                self._emit(
                    "tool.output",
                    {"call_id": call_id, "name": tool_name, "output": output_str},
                )

            if has_success:
                consecutive_errors = 0
            elif tool_calls:
                consecutive_errors += 1

            # Prepare next request with updated input
            request_kwargs["input"] = list(input_items)

        raise RunnerMaxIterationsError(
            f"Max iterations reached ({self._max_iters}). The model kept calling tools."
        )

    def _normalize_input(
        self, prompt_or_input: Union[str, Sequence[InputItem]]
    ) -> InputList:
        if isinstance(prompt_or_input, str):
            return [{"role": "user", "content": prompt_or_input}]
        return [dict(x) for x in prompt_or_input]

    def _emit(self, event: str, payload: Dict[str, Any]) -> None:
        if self._on_event:
            self._on_event(event, payload)


def _to_dict_items(items: List[Any]) -> List[Dict[str, Any]]:
    """Convert SDK items to plain dicts so we can round-trip them as input."""
    out: List[Dict[str, Any]] = []
    for it in items:
        if isinstance(it, dict):
            out.append(it)
        else:
            # Most OpenAI SDK objects support model_dump()
            dump = getattr(it, "model_dump", None)
            out.append(dump() if callable(dump) else dict(it))  # type: ignore[arg-type]
    return out


def run(
    client: Any,
    model: str,
    registry: ToolRegistry,
    prompt: str,
    *,
    tool_choice: Optional[Any] = None,
    instructions: Optional[str] = None,
    max_iters: int = 8,
    max_retry: int = 0,
) -> RunnerResult:
    """Convenience wrapper for one-off calls."""
    runner = ToolRunner(
        client, model=model, registry=registry, max_iters=max_iters, max_retry=max_retry
    )
    return runner.run(prompt, tool_choice=tool_choice, instructions=instructions)
