from __future__ import annotations

# standard
import json
import time
import inspect
import asyncio
from collections.abc import AsyncIterator

# third party
from sunwaee.core.logger import get_logger

# custom
from .engine.base import BaseEngine
from .engine.types import Message, Response, Role, Tool, ToolCall

_log = get_logger("sun.agent")


async def run(
    messages: list[Message],
    tools: list[Tool],
    engine: BaseEngine,
    max_iterations: int = 10,
) -> tuple[str, list[Message]]:
    """Run the ReAct loop.

    Returns:
        (final_text, new_messages) where new_messages are all messages
        appended during this run (not including the input messages).
        Callers can persist new_messages to a session file.
    """

    history = list(messages)
    start_len = len(history)
    tools_by_name = {t.name: t for t in tools}

    for iteration in range(max_iterations):
        _log.debug("Agent iteration %d/%d", iteration + 1, max_iterations)
        response = await engine.chat(history, tools)

        history.append(
            Message(
                role=Role.ASSISTANT,
                content=response.content,
                tool_calls=response.tool_calls,
                reasoning_content=response.reasoning_content,
                reasoning_signature=response.reasoning_signature,
            )
        )

        if (
            response.stop_reason and response.stop_reason.value != "tool_use"
        ) or not response.tool_calls:
            return response.content or "", history[start_len:]

        # execute all tool calls concurrently.
        results = await asyncio.gather(
            *[_execute(tc, tools_by_name) for tc in response.tool_calls]
        )

        for tc, result in zip(response.tool_calls, results):
            _log.debug("Tool %r → %s", tc.name, result[:120])
            history.append(
                Message(
                    role=Role.TOOL,
                    content=result,
                    tool_call_id=tc.id,
                )
            )

    _log.warning("Reached max_iterations (%d)", max_iterations)
    for msg in reversed(history):
        if msg.role.value == "assistant" and msg.content:  # pragma: no cover
            return msg.content, history[start_len:]
    return "Max iterations reached without a final response.", history[start_len:]


async def stream_run(
    messages: list[Message],
    tools: list[Tool],
    engine: BaseEngine,
    max_iterations: int = 10,
    new_messages: list[Message] | None = None,
) -> AsyncIterator[Response]:
    """Stream the ReAct loop, yielding Response chunks as they arrive.

    Text and reasoning chunks are yielded live. When the model stops with
    tool_use, tools are executed concurrently and the next iteration begins
    streaming immediately — transparent to the caller.

    Args:
        new_messages: optional out-parameter list; populated with all messages
            appended during this run (same as run()'s second return value).
    """

    history = list(messages)
    tools_by_name = {t.name: t for t in tools}
    out = new_messages if new_messages is not None else []

    for iteration in range(max_iterations):
        _log.debug("Agent stream iteration %d/%d", iteration + 1, max_iterations)

        accumulated_content = ""
        accumulated_reasoning = ""
        accumulated_signature: str | None = None
        accumulated_tool_calls: list[ToolCall] = []
        final_chunk: Response | None = None

        async for chunk in engine.stream(history, tools):
            yield chunk
            if chunk.synthetic:
                continue
            if chunk.content:
                accumulated_content += chunk.content
            if chunk.reasoning_content:
                accumulated_reasoning += chunk.reasoning_content
            if chunk.reasoning_signature:
                accumulated_signature = chunk.reasoning_signature
            if chunk.tool_calls:
                accumulated_tool_calls.extend(chunk.tool_calls)
            if chunk.stop_reason is not None:
                final_chunk = chunk

        assistant_msg = Message(
            role=Role.ASSISTANT,
            content=accumulated_content or None,
            reasoning_content=accumulated_reasoning or None,
            reasoning_signature=accumulated_signature,
            tool_calls=accumulated_tool_calls or None,
        )
        history.append(assistant_msg)
        out.append(assistant_msg)

        stop_reason = final_chunk.stop_reason if final_chunk else None
        if (
            stop_reason is None
            or stop_reason.value != "tool_use"
            or not accumulated_tool_calls
        ):
            return

        results = await asyncio.gather(
            *[_execute(tc, tools_by_name) for tc in accumulated_tool_calls]
        )

        for tc, result in zip(accumulated_tool_calls, results):
            _log.debug("Tool %r → %s", tc.name, result[:120])
            tool_msg = Message(role=Role.TOOL, content=result, tool_call_id=tc.id)
            history.append(tool_msg)
            out.append(tool_msg)

        # Signal to the caller that this batch of tool calls has completed.
        yield Response(synthetic=True, tool_calls=accumulated_tool_calls)

    _log.warning("Reached max_iterations (%d) in stream_run", max_iterations)


async def _execute(tc: ToolCall, tools_by_name: dict[str, Tool]) -> str:
    """Execute a single tool call, populate tc metadata, and return a JSON string result."""

    t_start = time.monotonic()

    tool = tools_by_name.get(tc.name)
    if tool is None:
        tc.error = f"Unknown tool: {tc.name!r}"
        tc.duration = 0.0
        return json.dumps({"ok": False, "error": tc.error})
    if tool.fn is None:  # pragma: no cover
        tc.error = f"Tool {tc.name!r} has no callable."
        tc.duration = 0.0
        return json.dumps({"ok": False, "error": tc.error})
    try:
        if inspect.iscoroutinefunction(tool.fn):
            raw = await tool.fn(**tc.arguments)
        else:
            raw = await asyncio.to_thread(tool.fn, **tc.arguments)
        if not isinstance(raw, str):
            raw = json.dumps(raw, default=str)  # pragma: no cover
        tc.duration = round(time.monotonic() - t_start, 3)
        try:
            parsed = json.loads(raw)
            if isinstance(parsed, dict):
                tc.results = [parsed]
            elif isinstance(parsed, list):  # pragma: no cover
                tc.results = [r for r in parsed if isinstance(r, dict)]
            else:  # pragma: no cover
                tc.results = [{"value": parsed}]
        except (json.JSONDecodeError, TypeError):  # pragma: no cover
            tc.results = [{"raw": raw}]
        return raw
    except Exception as exc:
        tc.duration = round(time.monotonic() - t_start, 3)
        tc.error = str(exc)
        _log.exception("Tool %r raised: %s", tc.name, exc)
        return json.dumps({"ok": False, "error": str(exc)})
