"""Agent loop — the core runtime. Model decides next step, executes tools, continues.

This is the Claude Code / Manus AI pattern: single-threaded agent loop where the model
IS the planner, executor, and evaluator — all in one loop.
"""

import inspect
from typing import Callable

import structlog
from pydantic import BaseModel, Field

from fliiq.runtime.agent.config import AgentConfig
from fliiq.runtime.agent.tools import ToolRegistry
from fliiq.runtime.llm.providers import BaseLLM, LLMResponse

log = structlog.get_logger()


class CancellationToken:
    """Cooperative cancellation — checked between iterations and tool calls."""

    def __init__(self):
        self._cancelled = False

    def cancel(self) -> None:
        self._cancelled = True

    @property
    def is_cancelled(self) -> bool:
        return self._cancelled


class AgentResult(BaseModel):
    """Result from an agent loop run."""

    messages: list[dict] = Field(default_factory=list)
    final_text: str = ""
    iterations: int = 0
    stop_reason: str = "end_turn"
    input_tokens: int = 0
    output_tokens: int = 0


def _build_assistant_message(response: LLMResponse) -> dict:
    """Build an assistant message from the LLM response, using raw content blocks."""
    if response.raw_content:
        return {"role": "assistant", "content": response.raw_content}
    return {"role": "assistant", "content": response.content}


def _build_tool_results_message(tool_results: list[dict]) -> dict:
    """Build a user message with tool results in Anthropic content block format."""
    return {
        "role": "user",
        "content": [
            {
                "type": "tool_result",
                "tool_use_id": result["tool_use_id"],
                "content": result["content"],
            }
            for result in tool_results
        ],
    }


def _save_audit(messages: list[dict], stop_reason: str, final_text: str = "") -> None:
    """Non-fatal audit save. Extract prompt from first user message."""
    try:
        from fliiq.runtime.agent.audit import extract_audit_trail, save_audit_trail

        prompt = ""
        for msg in messages:
            if msg.get("role") == "user":
                content = msg.get("content", "")
                prompt = content if isinstance(content, str) else str(content)[:200]
                break
        trail = extract_audit_trail(prompt, messages, stop_reason=stop_reason)
        trail.final_text = final_text
        save_audit_trail(trail)
    except Exception:
        pass


async def agent_loop(
    llm: BaseLLM,
    messages: list[dict],
    tools: ToolRegistry,
    config: AgentConfig | None = None,
    system: str | None = None,
    on_tool_call: Callable | None = None,
    on_before_tool_call: Callable[[str, dict], bool] | None = None,
    cancellation: CancellationToken | None = None,
    tool_defs_override: list[dict] | None = None,
    tool_skip_message: str = "Tool skipped by user",
) -> AgentResult:
    """Run the agent loop until the model stops or max iterations reached.

    Args:
        llm: LLM provider instance.
        messages: Initial message history (at minimum, a user message).
        tools: Tool registry with definitions and handlers.
        config: Agent configuration (mode, max_iterations).
        system: System prompt.
        on_tool_call: Optional callback(name, input, result) for each tool call.
        on_before_tool_call: Optional callback(name, input) -> bool. If returns False, tool is skipped.
        cancellation: Optional token for cooperative cancellation.
        tool_defs_override: Optional list of tool definitions to send to the LLM instead of
            the full registry. Allows callers to filter visible tools without mutating the registry.

    Returns:
        AgentResult with final messages, text, iteration count, and stop reason.
    """
    config = config or AgentConfig()
    iteration = 0
    total_input_tokens = 0
    total_output_tokens = 0

    while iteration < config.max_iterations:
        # Refresh tool definitions each iteration when no override is set.
        # This enables hot-loaded skills (install_skill) to become visible mid-run.
        if tool_defs_override is not None:
            tool_defs = tool_defs_override
        else:
            tool_defs = tools.get_tool_definitions()

        # Check cancellation at top of loop
        if cancellation and cancellation.is_cancelled:
            log.info("agent_loop_cancelled", iteration=iteration)
            return AgentResult(
                messages=messages,
                final_text="[Cancelled by user]",
                iterations=iteration,
                stop_reason="cancelled",
                input_tokens=total_input_tokens,
                output_tokens=total_output_tokens,
            )

        iteration += 1
        log.debug("agent_loop_iteration", iteration=iteration, max=config.max_iterations)

        response = await llm.generate(
            messages,
            system=system,
            tools=tool_defs if tool_defs else None,
        )
        if response.usage:
            total_input_tokens += response.usage.get("input_tokens", 0)
            total_output_tokens += response.usage.get("output_tokens", 0)

        # Append assistant message to history
        assistant_msg = _build_assistant_message(response)
        messages.append(assistant_msg)

        # If the model is done talking (no tool calls), we're finished
        if response.stop_reason == "end_turn" or not response.tool_calls:
            # Strip empty assistant message — avoids API 400 if conversation continues
            if not assistant_msg.get("content"):
                messages.pop()
            _save_audit(messages, response.stop_reason, response.content)
            return AgentResult(
                messages=messages,
                final_text=response.content,
                iterations=iteration,
                stop_reason=response.stop_reason,
                input_tokens=total_input_tokens,
                output_tokens=total_output_tokens,
            )

        # Execute tool calls and collect results
        tool_results = []
        for tc in response.tool_calls:
            # Check cancellation between tool calls
            if cancellation and cancellation.is_cancelled:
                tool_results.append({
                    "tool_use_id": tc.id,
                    "content": "Cancelled by user",
                })
                continue

            log.debug("agent_tool_call", tool=tc.name, input_keys=list(tc.input.keys()))

            # Check with on_before_tool_call callback (supervised mode)
            # Supports both sync and async callbacks for TUI compatibility
            if on_before_tool_call:
                approval = on_before_tool_call(tc.name, tc.input)
                if inspect.isawaitable(approval):
                    approval = await approval
            else:
                approval = True

            if not approval:
                result = tool_skip_message
            else:
                result = await tools.execute(tc.name, tc.input)

            if on_tool_call:
                _cb_result = on_tool_call(tc.name, tc.input, result)
                if inspect.isawaitable(_cb_result):
                    await _cb_result

            tool_results.append({
                "tool_use_id": tc.id,
                "content": result,
            })

        # Append tool results as a user message
        messages.append(_build_tool_results_message(tool_results))

        # Check cancellation after tool execution — return with results appended
        if cancellation and cancellation.is_cancelled:
            log.info("agent_loop_cancelled_after_tools", iteration=iteration)
            return AgentResult(
                messages=messages,
                final_text="[Cancelled by user]",
                iterations=iteration,
                stop_reason="cancelled",
                input_tokens=total_input_tokens,
                output_tokens=total_output_tokens,
            )

    # Reached max iterations
    log.warning("agent_loop_max_iterations", iterations=iteration)
    _save_audit(messages, "max_iterations", response.content if response else "")
    return AgentResult(
        messages=messages,
        final_text=response.content if response else "",
        iterations=iteration,
        stop_reason="max_iterations",
        input_tokens=total_input_tokens,
        output_tokens=total_output_tokens,
    )


def plan_was_approved(result: AgentResult) -> bool:
    """Check if a plan_complete tool returned 'approved' in the agent result."""
    for msg in reversed(result.messages):
        if msg.get("role") != "user":
            continue
        content = msg.get("content")
        if not isinstance(content, list):
            continue
        for block in content:
            if block.get("type") == "tool_result" and block.get("content") == "approved":
                return True
    return False
