"""ClawAgents ReAct Agent Loop

Single-loop ReAct executor inspired by deepagents/openclaw architecture.
Eliminates the separate Understand/Verify phases that added 2 unnecessary
LLM round-trips per iteration.

Flow: LLM → tool calls → LLM → tool calls → ... → final text answer

Robustness features retained:
  - Tool loop detection
  - Context-window guard with auto-compaction
  - Parallel tool execution
  - Tool-output truncation
  - Structured event callbacks (on_event)

Efficiency features (learned from deepagents/openclaw):
  - Adaptive token estimation multiplier (auto-calibrates after overflow)
  - Tool argument truncation in older messages (saves tokens)
  - Single-pass message filtering
"""

from __future__ import annotations

import asyncio
import json
import logging
import math
import re
import signal
import sys
import time
from dataclasses import dataclass
from typing import Any, Callable, Literal, Optional

from clawagents.providers.llm import LLMProvider, LLMMessage, LLMResponse
from clawagents.tools.registry import ToolRegistry, ParsedToolCall

logger = logging.getLogger(__name__)

AgentStatus = Literal["running", "done", "error"]

EventKind = Literal[
    "tool_call",
    "tool_result",
    "retry",
    "agent_done",
    "warn",
    "error",
    "context",
    "final_content",
    "approval_required",
    "tool_skipped",
]

OnEvent = Callable[[EventKind, dict[str, Any]], None]

# Hook types for extensibility without middleware overhead
BeforeLLMHook = Callable[[list["LLMMessage"]], list["LLMMessage"]]
BeforeToolHook = Callable[[str, dict[str, Any]], bool]
AfterToolHook = Callable[[str, dict[str, Any], "ToolResult"], "ToolResult"]


def _default_on_event(kind: EventKind, data: dict[str, Any]) -> None:
    """Default event handler: write to stderr (CLI mode)."""
    if kind == "tool_call":
        sys.stderr.write(f"\U0001f527 {data['name']}\n")
    elif kind == "retry":
        sys.stderr.write(f"[retry] {data['reason']}\n")
    elif kind == "agent_done":
        sys.stderr.write(
            f"\n\u2713 {data['tool_calls']} tool calls"
            f" \u00b7 {data['iterations']} iterations"
            f" \u00b7 {data['elapsed']:.1f}s\n"
        )
    elif kind == "final_content":
        sys.stdout.write(data["content"])
        sys.stdout.write("\n")
        sys.stdout.flush()
    elif kind == "warn":
        sys.stderr.write(f"[warn] {data['message']}\n")
    elif kind == "error":
        sys.stderr.write(f"[error] {data['phase']}: {data['message']}\n")
    elif kind == "context":
        sys.stderr.write(f"[context] {data['message']}\n")
    sys.stderr.flush()


@dataclass
class AgentState:
    messages: list[LLMMessage]
    current_task: str
    status: AgentStatus
    result: str
    iterations: int
    max_iterations: int
    tool_calls: int


BASE_SYSTEM_PROMPT = """You are a ClawAgent, an AI assistant that helps users accomplish tasks using tools. You respond with text and tool calls.

## Core Behavior
- Be concise and direct. Don't over-explain unless asked.
- NEVER add unnecessary preamble ("Sure!", "Great question!", "I'll now...").
- If the request is ambiguous, ask questions before acting.

## Doing Tasks
When the user asks you to do something:
1. Think briefly about your approach, then act immediately using tools.
2. After getting tool results, continue using more tools or provide the final answer.
3. When done, provide the final answer directly. Do NOT ask if the user wants more.

Keep working until the task is fully complete."""


# ─── Adaptive Token Estimation (learned from deepagents) ──────────────────

_CHARS_PER_TOKEN = 4


def _estimate_tokens(text: str, multiplier: float = 1.0) -> int:
    return math.ceil((len(text) / _CHARS_PER_TOKEN) * multiplier)


def _estimate_messages_tokens(messages: list[LLMMessage], multiplier: float = 1.0) -> int:
    return sum(_estimate_tokens(m.content, multiplier) for m in messages)


# ─── Tool Argument Truncation in Old Messages (learned from deepagents) ───

_MAX_ARG_LENGTH = 2000
_ARG_TRUNCATION_MARKER = "...(argument truncated)"
_RECENT_PROTECTED_COUNT = 20
_TRUNCATABLE_RE = re.compile(
    r'\{"tool":\s*"(write_file|edit_file|create_file)".*?"args":\s*\{'
)


def _truncate_old_tool_args(
    messages: list[LLMMessage], protect_recent: int = _RECENT_PROTECTED_COUNT,
) -> list[LLMMessage]:
    if len(messages) <= protect_recent:
        return messages

    cutoff = len(messages) - protect_recent
    result: list[LLMMessage] = []

    for i, m in enumerate(messages):
        if (
            i < cutoff
            and m.role == "assistant"
            and _TRUNCATABLE_RE.search(m.content)
            and len(m.content) > _MAX_ARG_LENGTH
        ):
            result.append(LLMMessage(
                role=m.role,
                content=m.content[:_MAX_ARG_LENGTH] + _ARG_TRUNCATION_MARKER,
            ))
        else:
            result.append(m)

    return result


# ─── Tool Loop Detection ──────────────────────────────────────────────────


class _ToolCallTracker:
    def __init__(self, window_size: int = 12, max_repeats: int = 3):
        self._history: list[str] = []
        self._window_size = window_size
        self._max_repeats = max_repeats

    def _key(self, tool_name: str, args: dict) -> str:
        return f"{tool_name}:{json.dumps(args, sort_keys=True)}"

    def record(self, tool_name: str, args: dict) -> None:
        self._history.append(self._key(tool_name, args))
        if len(self._history) > self._window_size:
            self._history.pop(0)

    def is_looping(self, tool_name: str, args: dict) -> bool:
        key = self._key(tool_name, args)
        return self._history.count(key) >= self._max_repeats

    def is_looping_batch(self, calls: list[ParsedToolCall]) -> bool:
        return any(self.is_looping(c.tool_name, c.args) for c in calls)

    def record_batch(self, calls: list[ParsedToolCall]) -> None:
        for c in calls:
            self.record(c.tool_name, c.args)


# ─── Context Window Guard with Auto-Compaction ────────────────────────────

_CONTEXT_BUDGET_RATIO = 0.75
_RECENT_MESSAGES_TO_KEEP = 6


async def _compact_if_needed(
    messages: list[LLMMessage],
    context_window: int,
    llm: LLMProvider,
    emit: OnEvent,
    token_multiplier: float = 1.0,
) -> list[LLMMessage]:
    # Phase 1: truncate tool args in older messages (cheap, no LLM call)
    messages = _truncate_old_tool_args(messages)

    budget = int(context_window * _CONTEXT_BUDGET_RATIO)
    current_tokens = _estimate_messages_tokens(messages, token_multiplier)

    if current_tokens <= budget:
        return messages

    emit("context", {"message": f"~{current_tokens} tokens exceeds budget {budget} — compacting"})

    # Single-pass split
    system_msgs: list[LLMMessage] = []
    non_system: list[LLMMessage] = []
    for m in messages:
        (system_msgs if m.role == "system" else non_system).append(m)

    if len(non_system) <= _RECENT_MESSAGES_TO_KEEP:
        return messages

    recent_count = min(_RECENT_MESSAGES_TO_KEEP, len(non_system))
    older = non_system[:-recent_count]
    recent = non_system[-recent_count:]

    # Phase 8: Offload full history to file before summarizing
    offload_path = _offload_history(older)
    if offload_path:
        emit("context", {"message": f"offloaded {len(older)} messages to {offload_path}"})

    text_log = "\n\n".join(
        f"[{m.role.upper()}]: {m.content}" for m in older
    )

    summary_prompt = (
        "Compress the following agent conversation history into a concise summary. "
        "Keep key facts, file paths, errors, and tool results. Be brief.\n\n"
        + text_log
    )

    try:
        resp = await llm.chat([LLMMessage(role="user", content=summary_prompt)])
        summary = LLMMessage(
            role="assistant",
            content=f"[Compacted History] {resp.content}",
        )
        emit("context", {"message": f"compacted {len(older)} messages into summary"})
        return [*system_msgs, summary, *recent]
    except Exception:
        logger.debug("Compaction LLM call failed", exc_info=True)
        emit("context", {"message": "compaction failed — dropping oldest messages"})
        return [*system_msgs, *recent]


# ─── History Offloading ───────────────────────────────────────────────────

from pathlib import Path

_HISTORY_DIR = Path.cwd() / ".clawagents" / "history"


def _offload_history(messages: list[LLMMessage]) -> str | None:
    """Save older messages to a JSON file before compaction."""
    try:
        _HISTORY_DIR.mkdir(parents=True, exist_ok=True)
        ts = int(time.time())
        path = _HISTORY_DIR / f"compacted_{ts}_{len(messages)}msgs.json"
        data = [{"role": m.role, "content": m.content} for m in messages]
        path.write_text(json.dumps(data, indent=2), "utf-8")
        return str(path)
    except Exception:
        logger.debug("History offload failed", exc_info=True)
        return None


# ─── Helpers ──────────────────────────────────────────────────────────────


def _make_buffer():
    buf: list[str] = []
    def on_chunk(chunk: str) -> None:
        buf.append(chunk)
    return buf, on_chunk


# ─── ReAct Loop ──────────────────────────────────────────────────────────

MAX_TOOL_ROUNDS = 15


async def run_agent_graph(
    task: str,
    llm: LLMProvider,
    tools: Optional[ToolRegistry] = None,
    system_prompt: Optional[str] = None,
    max_iterations: int = 3,
    streaming: bool = True,
    context_window: int = 128_000,
    on_event: Optional[OnEvent] = None,
    before_llm: Optional[BeforeLLMHook] = None,
    before_tool: Optional[BeforeToolHook] = None,
    after_tool: Optional[AfterToolHook] = None,
) -> AgentState:
    """Single ReAct loop: LLM → tools → LLM → tools → ... → final answer."""
    registry = tools or ToolRegistry()
    tool_desc = registry.describe_for_llm()
    loop_tracker = _ToolCallTracker()
    emit = on_event or _default_on_event

    token_multiplier = 1.0

    prompt_to_use = system_prompt or BASE_SYSTEM_PROMPT

    messages: list[LLMMessage] = [
        LLMMessage(role="system", content=f"{prompt_to_use}\n\n{tool_desc}"),
        LLMMessage(role="user", content=task),
    ]

    state = AgentState(
        messages=messages,
        current_task=task,
        status="running",
        result="",
        iterations=0,
        max_iterations=max_iterations,
        tool_calls=0,
    )

    cancel_event = asyncio.Event()
    loop = asyncio.get_running_loop()

    def _on_sigint() -> None:
        emit("warn", {"message": "interrupted"})
        cancel_event.set()

    try:
        loop.add_signal_handler(signal.SIGINT, _on_sigint)
    except (NotImplementedError, OSError):
        pass

    t0 = time.monotonic()

    try:
        for round_idx in range(MAX_TOOL_ROUNDS):
            if cancel_event.is_set():
                state.status = "done"
                state.result = state.result or "[cancelled]"
                break

            messages = await _compact_if_needed(
                messages, context_window, llm, emit, token_multiplier,
            )

            # ── before_llm hook ──
            if before_llm:
                messages = before_llm(messages)

            buf, on_chunk = _make_buffer()
            try:
                response = await llm.chat(
                    messages,
                    on_chunk=on_chunk if streaming else None,
                    cancel_event=cancel_event,
                )
            except Exception as err:
                err_msg = str(err).lower()
                if "context" in err_msg or "token" in err_msg:
                    observed_ratio = context_window / max(
                        _estimate_messages_tokens(messages, 1.0), 1,
                    )
                    token_multiplier = min(observed_ratio * 1.1, 3.0)
                    emit("context", {
                        "message": f"token overflow — calibrated multiplier to {token_multiplier:.2f}",
                    })
                    messages = await _compact_if_needed(
                        messages, context_window, llm, emit, token_multiplier,
                    )
                    continue

                logger.exception("LLM call failed at round %d", round_idx)
                emit("error", {"phase": "llm_call", "message": str(err)})
                state.status = "error"
                state.result = str(err)
                break

            if response.partial and not response.content.strip():
                emit("warn", {"message": "interrupted — no content received"})
                state.status = "done"
                state.result = state.result or "[interrupted]"
                break

            tool_calls = registry.parse_tool_calls(response.content)

            if not tool_calls:
                state.result = response.content
                state.status = "done"
                state.iterations += 1
                emit("final_content", {"content": response.content})
                messages.append(LLMMessage(role="assistant", content=response.content))
                break

            if loop_tracker.is_looping_batch(tool_calls):
                names = ", ".join(c.tool_name for c in tool_calls)
                emit("warn", {"message": f"tool loop detected ({names}) — breaking"})
                state.status = "done"
                state.result = f"Tool loop detected ({names}). Stopping."
                state.iterations += 1
                break

            if len(tool_calls) == 1:
                call = tool_calls[0]
                emit("tool_call", {"name": call.tool_name})

                # ── before_tool hook ──
                if before_tool and not before_tool(call.tool_name, call.args):
                    emit("tool_skipped", {"name": call.tool_name, "reason": "rejected by before_tool hook"})
                    messages.append(LLMMessage(role="user", content=f"[Tool Skipped] {call.tool_name} was not approved."))
                    continue

                loop_tracker.record(call.tool_name, call.args)

                tool_result = await registry.execute_tool(call.tool_name, call.args)
                state.tool_calls += 1

                # ── after_tool hook ──
                if after_tool:
                    tool_result = after_tool(call.tool_name, call.args, tool_result)

                tool_output = (
                    tool_result.output if tool_result.success else f"Error: {tool_result.error}"
                )
                emit("tool_result", {
                    "name": call.tool_name,
                    "success": tool_result.success,
                    "preview": tool_output[:120],
                })

                messages.append(
                    LLMMessage(role="assistant", content=f'{{"tool": "{call.tool_name}", "args": {json.dumps(call.args)}}}')
                )
                messages.append(
                    LLMMessage(role="user", content=f"[Tool Result] {tool_output}")
                )
            else:
                # ── before_tool hook (parallel) — filter out rejected calls ──
                approved_calls = tool_calls
                if before_tool:
                    approved_calls = [c for c in tool_calls if before_tool(c.tool_name, c.args)]
                    skipped = [c for c in tool_calls if c not in approved_calls]
                    for c in skipped:
                        emit("tool_skipped", {"name": c.tool_name, "reason": "rejected by before_tool hook"})
                    if not approved_calls:
                        messages.append(LLMMessage(role="user", content="[Tool Skipped] All tool calls were not approved."))
                        continue

                for call in approved_calls:
                    emit("tool_call", {"name": call.tool_name})
                loop_tracker.record_batch(approved_calls)

                results = await registry.execute_tools_parallel(approved_calls)
                state.tool_calls += len(approved_calls)

                # ── after_tool hook (parallel) ──
                if after_tool:
                    results = [after_tool(c.tool_name, c.args, r) for c, r in zip(approved_calls, results)]

                call_summaries: list[str] = []
                for call, result in zip(approved_calls, results):
                    output = result.output if result.success else f"Error: {result.error}"
                    emit("tool_result", {
                        "name": call.tool_name,
                        "success": result.success,
                        "preview": output[:120],
                    })
                    call_summaries.append(f"{call.tool_name}({json.dumps(call.args)}) => {output}")

                tool_call_str = json.dumps([
                    {"tool": c.tool_name, "args": c.args} for c in approved_calls
                ])
                messages.append(
                    LLMMessage(role="assistant", content=tool_call_str)
                )
                messages.append(
                    LLMMessage(
                        role="user",
                        content="[Tool Results]\n" + "\n".join(call_summaries),
                    )
                )

        else:
            emit("warn", {"message": f"reached max {MAX_TOOL_ROUNDS} tool rounds"})
            state.status = "done"
            state.iterations += 1

    except KeyboardInterrupt:
        emit("warn", {"message": "interrupted"})
        state.status = "done"
        state.result = state.result or "[interrupted]"
    except asyncio.CancelledError:
        emit("warn", {"message": "cancelled"})
        state.status = "done"
        state.result = state.result or "[cancelled]"
    finally:
        try:
            loop.remove_signal_handler(signal.SIGINT)
        except (NotImplementedError, OSError):
            pass

    elapsed = time.monotonic() - t0
    state.messages = messages
    emit("agent_done", {
        "tool_calls": state.tool_calls,
        "iterations": state.iterations,
        "elapsed": elapsed,
    })
    return state
