# retoor <retoor@molodetz.nl>

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any, Callable, Optional

from ..config import Settings
from ..errors import LLMError
from ..chunks import reset_store, set_store
from ..cost import reset_tracker, set_tracker
from .compaction import compact_messages, context_size
from .state import AgentState, reset_state, set_state

logger = logging.getLogger("devii.agentic.loop")

TraceCallback = Callable[[str, str, str], None]
OUTPUT_CAP_CHARS = 400_000

REFLECTION_TRIGGER = (
    "[reflection-trigger] One or more tool calls returned an error. Call reflect() with the "
    "observation, root-cause conclusion, and next action before retrying. Do not repeat the "
    "same call without diagnosis."
)
PLAN_VIOLATION = (
    "Protocol violation: your first tool call must be plan(). Restart with a structured plan "
    "before any other action."
)
VERIFICATION_GATE = (
    "[verification-gate] You produced a final answer after performing changes without calling "
    "verify(). Confirm the change took effect and call verify() with a summary. If verification "
    "truly does not apply (for example the action only read data and changed nothing), do NOT "
    "reply with a bare disclaimer: give the user the full answer they asked for - including any "
    "rows or data you retrieved - and you may note 'No verification applicable: <reason>' at the "
    "end. Never drop the requested data."
)
ITERATION_LIMIT_MESSAGE = "[stopped] Maximum iterations reached without a final answer."


LABEL_KEYS = (
    "path",
    "post_slug",
    "project_slug",
    "gist_slug",
    "news_slug",
    "username",
    "comment_uid",
    "notification_uid",
    "attachment_uid",
    "poll_uid",
    "receiver_uid",
    "uid",
    "name",
    "q",
    "query",
    "title",
    "goal",
    "task",
    "email",
)
LABEL_MAX_CHARS = 60


def call_label(call: dict[str, Any]) -> str:
    raw = call.get("function", {}).get("arguments") or "{}"
    try:
        arguments = json.loads(raw) if isinstance(raw, str) else raw
    except (ValueError, TypeError):
        return ""
    if not isinstance(arguments, dict):
        return ""

    def clean(value: Any) -> str:
        return str(value).replace("\n", " ").strip()[:LABEL_MAX_CHARS]

    if arguments.get("url"):
        method = arguments.get("method")
        target = str(arguments["url"]).strip()
        return clean(f"{str(method).strip().upper()} {target}" if method else target)
    if arguments.get("target_uid"):
        target_type = arguments.get("target_type")
        target = clean(arguments["target_uid"])
        return f"{target_type}/{target}" if target_type else target
    if arguments.get("command"):
        return clean(arguments["command"])
    if arguments.get("instance"):
        instance = clean(arguments["instance"])
        action = arguments.get("action")
        return f"{clean(action)} {instance}" if action else instance
    for key in LABEL_KEYS:
        value = arguments.get(key)
        if value not in (None, ""):
            return clean(value)
    for value in arguments.values():
        if isinstance(value, (str, int, float, bool)) and clean(value):
            return clean(value)
    return ""


def normalize(message: dict[str, Any]) -> dict[str, Any]:
    normalized: dict[str, Any] = {
        "role": message.get("role", "assistant"),
        "content": message.get("content") or "",
    }
    if message.get("tool_calls"):
        normalized["tool_calls"] = message["tool_calls"]
    return normalized


def _is_error(result: str) -> bool:
    try:
        return (
            isinstance(json.loads(result), dict)
            and json.loads(result).get("error") is not None
        )
    except (ValueError, TypeError):
        return False


def _summary(result: str) -> str:
    try:
        parsed = json.loads(result)
    except (ValueError, TypeError):
        return ""
    if not isinstance(parsed, dict):
        return ""
    if parsed.get("error"):
        return str(parsed.get("message") or parsed.get("error"))
    if "status" in parsed:
        return str(parsed["status"])
    return ""


async def _run_tool_call(dispatcher: Any, call: dict[str, Any]) -> str:
    function = call.get("function", {})
    name = function.get("name", "")
    raw_arguments = function.get("arguments") or "{}"
    try:
        arguments = (
            json.loads(raw_arguments)
            if isinstance(raw_arguments, str)
            else raw_arguments
        )
    except json.JSONDecodeError as exc:
        return json.dumps(
            {
                "error": "tool_input_truncated",
                "message": (
                    f"The arguments for {name or 'this tool'} were cut off and could not be parsed "
                    f"({exc.msg} at position {exc.pos}); the model output hit its length limit. "
                    "Emit only one write tool call per turn (do not batch several file writes into a "
                    "single response) and resend this one call on its own."
                ),
            }
        )
    if not isinstance(arguments, dict):
        return json.dumps(
            {
                "error": "tool_input_error",
                "message": "Invalid arguments: arguments must be an object",
            }
        )
    return await dispatcher.dispatch(name, arguments)


async def react_loop(
    llm: Any,
    dispatcher: Any,
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]],
    state: AgentState,
    settings: Settings,
    max_iterations: int,
    plan_required: bool,
    verify_required: bool,
    on_trace: Optional[TraceCallback] = None,
    cost_tracker: Any = None,
    chunk_store: Any = None,
) -> str:
    def trace(event: str, name: str = "", detail: str = "") -> None:
        if on_trace is not None:
            on_trace(event, name, detail)

    token = set_state(state)
    cost_token = set_tracker(cost_tracker) if cost_tracker is not None else None
    chunk_token = set_store(chunk_store) if chunk_store is not None else None
    final_content = ""
    try:
        while state.iteration < max_iterations:
            state.iteration += 1

            if context_size(messages) > settings.context_compact_threshold:
                trace("compact")
                messages[:] = await compact_messages(
                    llm, messages, settings.context_keep_tail
                )

            try:
                message = await llm.complete(messages, tools)
            except LLMError as exc:
                logger.info("LLM error: %s", exc.message)
                return f"[model error] {exc.message}"

            messages.append(normalize(message))
            tool_calls = message.get("tool_calls") or []

            if tool_calls:
                needs_plan = (
                    plan_required
                    and state.plan is None
                    and any(
                        call["function"]["name"] != "plan"
                        and not dispatcher.is_read_only(call["function"]["name"])
                        for call in tool_calls
                    )
                )
                if needs_plan:
                    trace("plan-gate")
                    for call in tool_calls:
                        messages.append(
                            {
                                "role": "tool",
                                "tool_call_id": call.get("id", ""),
                                "name": call.get("function", {}).get("name", ""),
                                "content": json.dumps(
                                    {"error": "protocol", "message": PLAN_VIOLATION}
                                ),
                            }
                        )
                    continue

                for call in tool_calls:
                    trace(
                        "call",
                        call.get("function", {}).get("name", ""),
                        call_label(call),
                    )

                results = await asyncio.gather(
                    *[_run_tool_call(dispatcher, call) for call in tool_calls]
                )

                any_error = False
                for call, result in zip(tool_calls, results):
                    if len(result) > OUTPUT_CAP_CHARS:
                        result = (
                            result[:OUTPUT_CAP_CHARS]
                            + f"\n...[truncated {len(result)} chars]"
                        )
                    name = call.get("function", {}).get("name", "")
                    messages.append(
                        {
                            "role": "tool",
                            "tool_call_id": call.get("id", ""),
                            "name": name,
                            "content": result,
                        }
                    )
                    label = call_label(call)
                    summary = _summary(result)
                    detail = " ".join(
                        part
                        for part in (label, f"({summary})" if summary else "")
                        if part
                    )
                    if _is_error(result):
                        any_error = True
                        trace("err", name, detail)
                    else:
                        trace("ok", name, detail)

                if any_error:
                    trace("reflect-trigger")
                    messages.append({"role": "user", "content": REFLECTION_TRIGGER})
                continue

            content = message.get("content")
            if content:
                if (
                    verify_required
                    and not state.verified
                    and not state.gate_triggered
                    and state.mutations
                ):
                    state.gate_triggered = True
                    trace("verify-gate")
                    messages.append({"role": "user", "content": VERIFICATION_GATE})
                    continue
                final_content = content
                break

        if state.iteration >= max_iterations and not final_content:
            return ITERATION_LIMIT_MESSAGE
        return final_content
    finally:
        if chunk_token is not None:
            reset_store(chunk_token)
        if cost_token is not None:
            reset_tracker(cost_token)
        reset_state(token)
