"""Toolformer reasoning pattern implementation.

Toolformer: Self-supervised tool usage learning
- The model learns to call APIs by generating [API_CALL(...)] tokens inline
- After the API call, the model observes the result and continues generation
- This pattern teaches the model when and how to appropriately use tools
"""

from __future__ import annotations

import logging
import re
from typing import Any

from hermes.core.types import AgentStrategy

logger = logging.getLogger(__name__)


class ToolformerReasoner:
    """Implements the Toolformer pattern for self-supervised tool use."""

    def __init__(self) -> None:
        self.strategy = AgentStrategy.TOOLFORMER

    def create_prompt(
        self, task: str, tools: list[dict[str, Any]], history: list[dict[str, str]] | None = None
    ) -> str:
        """Create Toolformer prompt with inline API call format."""
        tool_descriptions = "\n".join(
            f"- {t['name']}: {t['description']}" for t in tools
        )

        history_text = ""
        if history:
            history_text = "\nPrevious context:\n"
            for h in history:
                if "text" in h:
                    history_text += f"{h['text']}\n"

        return f"""You are an AI agent that uses tools by making inline API calls.

Task: {task}

Available tools:
{tool_descriptions}
{history_text}
To use a tool, embed an API call directly in your text using this format:

[API_CALL: tool_name(arg1="value1", arg2="value2")]

The tool result will be provided, and you can continue your response.

Example:
I need to find the latest news about AI.
[API_CALL: search_web(query="latest AI news 2026")]
Now I have the search results, I can answer the question: ...

Rules:
- Only call one API at a time
- Wait for the result before continuing
- Use the tool output to inform your response
- If you don't need any tools, just answer directly"""
    def parse_response(self, response: str) -> dict[str, Any]:
        """Parse Toolformer response, extracting API calls."""
        result: dict[str, Any] = {
            "text": response,
            "api_calls": [],
            "has_api_call": False,
            "final_answer": None,
        }

        api_call_pattern = re.compile(
            r'\[API_CALL:\s*(\w+)\s*\(([^)]*)\)\s*\]'
        )

        matches = api_call_pattern.findall(response)
        for tool_name, args_str in matches:
            try:
                args = self._parse_args(args_str)
                result["api_calls"].append({
                    "tool": tool_name,
                    "arguments": args,
                })
            except Exception as e:
                logger.warning(f"Failed to parse API call: {e}")

        result["has_api_call"] = len(result["api_calls"]) > 0

        if not result["has_api_call"]:
            result["final_answer"] = response.strip()

        return result

    def create_observation_prompt(
        self, original: str, api_call: dict[str, Any], observation: str
    ) -> str:
        """Create prompt to continue generation after tool observation."""
        return f"""Continue your response after receiving the tool result.

Your original text:
{original}

You called:
[API_CALL: {api_call['tool']}({api_call['arguments']})]

Tool result:
{observation}

Now incorporate this result into your response and continue.
If you need more information, you can make another API call.
If you have enough information, provide your final answer."""
    def parse_observation_response(self, response: str) -> dict[str, Any]:
        """Parse the continued response after an observation."""
        result: dict[str, Any] = {"text": response, "has_more_calls": False, "final_answer": None}

        if "[API_CALL:" in response:
            result["has_more_calls"] = True
            parsed = self.parse_response(response)
            result["api_calls"] = parsed.get("api_calls", [])
        else:
            result["final_answer"] = response.strip()

        return result

    def should_continue(self, parsed: dict[str, Any], max_calls: int, call_count: int) -> bool:
        """Determine if the agent should make more API calls."""
        if call_count >= max_calls:
            return False
        if parsed.get("final_answer"):
            return False
        return not (not parsed.get("has_api_call") and not parsed.get("has_more_calls"))

    def _parse_args(self, args_str: str) -> dict[str, Any]:
        """Parse argument string into a dictionary."""
        args: dict[str, Any] = {}
        if not args_str.strip():
            return args

        parts = re.findall(r'(\w+)\s*=\s*("[^"]*"|\'[^\']*\'|\S+)', args_str)
        for key, value in parts:
            cleaned = value.strip('"').strip("'")
            args[key] = cleaned

        return args