"""
Multi-step agent executor implementing the ReAct pattern.

ReAct (Reason + Act) interleaves thinking and tool use: the model
reasons about what to do, calls a tool, observes the result, then
decides whether to call another tool or give a final answer.

The loop has hard guardrails for a small model:

* *max_steps* (default 5) prevents infinite loops.
* If the model never produces ``<answer>``, the last response is used.
* Tool errors are fed back as observations so the model can recover.
* Repeated identical tool calls break the loop early.

The executor is pure Python with no Streamlit dependency. It takes
a ``GenerationEngine`` and ``ToolRegistry``, and returns an
``AgentResult``.
"""

from __future__ import annotations

from typing import Iterator

from src.agents.parser import coerce_param, parse_tool_call
from src.agents.registry import ToolRegistry
from src.agents.schemas import AgentResult, AgentStep, ToolCall

# Optional guardrails import — available when src.guardrails is installed.
try:
    from src.guardrails.middleware import GuardrailsMiddleware
except ImportError:  # pragma: no cover
    GuardrailsMiddleware = None  # type: ignore[assignment, misc]


# ---------------------------------------------------------------------------
# System prompt template
# ---------------------------------------------------------------------------

_SYSTEM_TEMPLATE = """\
You are a helpful assistant with access to tools.

## How to use a tool

Thought: your reasoning about what to do
<tool>tool_name</tool>
<param name="param_name">value</param>

## How to give a final answer

Thought: your reasoning
<answer>your final answer here</answer>

## Rules
- Use exactly ONE tool per response.
- Always wrap your final answer in <answer></answer> tags.
- If you can answer without tools, use <answer> immediately.
- When the user asks about documents, files, policies, reports, revenue, or project-specific information, ALWAYS use retrieve_documents first — do not answer from your own knowledge or use web_search.
- Only use web_search for current events, live data, or topics clearly outside the uploaded documents.
- If unsure whether to use retrieve_documents or web_search, prefer retrieve_documents.

## Examples

User: What does the document say about pricing?
Assistant: Thought: The user is asking about document content, so I should search the uploaded documents.
<tool>retrieve_documents</tool>
<param name="query">pricing</param>

User: What models does the project support?
Assistant: Thought: This is about uploaded project documentation, so I should search the documents.
<tool>retrieve_documents</tool>
<param name="query">supported models</param>

User: What is the average salary?
Assistant: Thought: I need to query the data table for the average salary.
<tool>query_data</tool>
<param name="query">SELECT AVG(salary) FROM employees</param>

User: What is the latest news about Python?
Assistant: Thought: This is about current events, so I need to search the web.
<tool>web_search</tool>
<param name="query">latest Python news</param>

User: Hello!
Assistant: Thought: This is a greeting, no tools needed.
<answer>Hello! How can I help you today?</answer>

{tool_descriptions}"""


class AgentExecutor:
    """ReAct agent loop with tool calling."""

    def __init__(
        self,
        engine,
        registry: ToolRegistry,
        max_steps: int = 5,
        temperature: float = 0.3,
        max_tokens_per_step: int = 300,
        guardrails: GuardrailsMiddleware | None = None,
        extra_instructions: str = "",
    ) -> None:
        self.engine = engine
        self.registry = registry
        self.max_steps = max_steps
        self.temperature = temperature
        self.max_tokens_per_step = max_tokens_per_step
        self.guardrails = guardrails
        # Optional domain-specific rules appended to the system prompt
        # (e.g. financial period-grounding), kept out of the generic template.
        self.extra_instructions = extra_instructions

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def run(
        self,
        user_query: str,
        chat_history: list[dict] | None = None,
    ) -> AgentResult:
        """Execute the agent loop and return the final result."""
        # Scan user input.
        if self.guardrails:
            input_scan = self.guardrails.scan_input(user_query)
            if input_scan.blocked:
                return AgentResult(
                    answer="I can't process that request — it was flagged by safety filters.",
                    steps=[],
                    num_steps=0,
                    finished=True,
                    guardrail_flags=[f.__dict__ for f in input_scan.flags],
                )

        steps: list[AgentStep] = []
        system_prompt = self._build_system_prompt()
        prev_call: tuple[str, str] | None = None  # (tool_name, args_key) for repeat detection

        for _ in range(self.max_steps):
            messages = self._build_step_messages(
                system_prompt, user_query, steps, chat_history,
            )
            raw_output = self.engine.generate_answer(
                messages,
                max_tokens=self.max_tokens_per_step,
                temperature=self.temperature,
            )

            thought, tool_call, final_answer = parse_tool_call(raw_output)

            # Final answer path.
            if final_answer is not None:
                step = AgentStep(thought=thought)
                steps.append(step)
                result = AgentResult(
                    answer=final_answer,
                    steps=steps,
                    num_steps=len(steps),
                    finished=True,
                )
                return self._scan_answer(result)

            # Tool call path.
            assert tool_call is not None
            call_key = (tool_call.tool_name, str(sorted(tool_call.arguments.items())))

            # Detect repeated identical tool call.
            if call_key == prev_call:
                step = AgentStep(
                    thought=thought,
                    tool_call=tool_call,
                    error="Repeated tool call — stopping.",
                )
                steps.append(step)
                break
            prev_call = call_key

            observation = self._execute_tool(tool_call)
            step = AgentStep(
                thought=thought,
                tool_call=tool_call,
                observation=observation,
            )
            steps.append(step)

        # Exhausted max_steps — use last observation or raw output as answer.
        fallback = steps[-1].observation if steps else ""
        result = AgentResult(
            answer=fallback or "I was unable to determine an answer.",
            steps=steps,
            num_steps=len(steps),
            finished=False,
        )
        return self._scan_answer(result)

    def run_stream(
        self,
        user_query: str,
        chat_history: list[dict] | None = None,
    ) -> Iterator[tuple[int, str, str]]:
        """Yield ``(step_index, event_type, content)`` tuples for streaming UI.

        *event_type* is one of: ``"thought"``, ``"tool_call"``,
        ``"observation"``, ``"answer"``, ``"guardrail"``.
        """
        # Scan user input.
        if self.guardrails:
            input_scan = self.guardrails.scan_input(user_query)
            if input_scan.blocked:
                yield 0, "guardrail", "; ".join(f.description for f in input_scan.flags)
                yield 0, "answer", "I can't process that request — it was flagged by safety filters."
                return

        steps: list[AgentStep] = []
        system_prompt = self._build_system_prompt()
        prev_call: tuple[str, str] | None = None

        for step_idx in range(self.max_steps):
            messages = self._build_step_messages(
                system_prompt, user_query, steps, chat_history,
            )
            raw_output = self.engine.generate_answer(
                messages,
                max_tokens=self.max_tokens_per_step,
                temperature=self.temperature,
            )

            thought, tool_call, final_answer = parse_tool_call(raw_output)

            if thought:
                yield step_idx, "thought", thought

            # Final answer.
            if final_answer is not None:
                steps.append(AgentStep(thought=thought))
                if self.guardrails:
                    out_scan = self.guardrails.scan_output(final_answer)
                    if out_scan.flags:
                        yield step_idx, "guardrail", "; ".join(f.description for f in out_scan.flags)
                    if out_scan.blocked:
                        yield step_idx, "answer", "The response was blocked by safety filters."
                        return
                yield step_idx, "answer", final_answer
                return

            # Tool call.
            assert tool_call is not None
            call_key = (tool_call.tool_name, str(sorted(tool_call.arguments.items())))

            call_desc = "{}({})".format(
                tool_call.tool_name,
                ", ".join("{}={}".format(k, v) for k, v in tool_call.arguments.items()),
            )
            yield step_idx, "tool_call", call_desc

            if call_key == prev_call:
                steps.append(AgentStep(
                    thought=thought, tool_call=tool_call,
                    error="Repeated tool call — stopping.",
                ))
                break
            prev_call = call_key

            observation = self._execute_tool(tool_call)
            steps.append(AgentStep(
                thought=thought, tool_call=tool_call, observation=observation,
            ))
            yield step_idx, "observation", observation

        # Exhausted steps — yield whatever we have.
        fallback = steps[-1].observation if steps else "I was unable to determine an answer."
        yield len(steps) - 1, "answer", fallback

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _scan_answer(self, result: AgentResult) -> AgentResult:
        """Run guardrail checks on the final answer, if configured."""
        if not self.guardrails:
            return result
        out_scan = self.guardrails.scan_output(result.answer)
        if out_scan.flags:
            result.guardrail_flags = [
                {"check_name": f.check_name, "severity": f.severity.value, "description": f.description}
                for f in out_scan.flags
            ]
        if out_scan.blocked:
            result.answer = "The response was blocked by safety filters."
        return result

    def _build_system_prompt(self) -> str:
        prompt = _SYSTEM_TEMPLATE.format(
            tool_descriptions=self.registry.format_tool_descriptions(),
        )
        if self.extra_instructions:
            prompt += "\n\n## Additional rules\n" + self.extra_instructions
        return prompt

    def _build_step_messages(
        self,
        system_prompt: str,
        user_query: str,
        steps: list[AgentStep],
        chat_history: list[dict] | None,
    ) -> list[dict]:
        messages: list[dict] = [{"role": "system", "content": system_prompt}]

        # Include recent chat history (last 4 turns for context window budget).
        if chat_history:
            messages.extend(chat_history[-4:])

        # Append /no_think for Qwen3 to suppress <think> blocks — they waste
        # the token budget on internal reasoning that duplicates our Thought: line.
        # Other architectures (Gemma3, Gemma4) don't use this control token.
        arch = getattr(self.engine, "arch", "qwen3")
        suffix = " /no_think" if arch == "qwen3" else ""
        messages.append({"role": "user", "content": user_query + suffix})

        # Encode previous steps as assistant/user turn pairs.
        for step in steps:
            # Reconstruct the assistant's raw output.
            raw = step.tool_call.raw_text if step.tool_call and step.tool_call.raw_text else ""
            if not raw and step.thought:
                raw = "Thought: {}".format(step.thought)
            if raw:
                messages.append({"role": "assistant", "content": raw})
            if step.observation:
                messages.append({
                    "role": "user",
                    "content": "Observation: {}".format(step.observation),
                })

        return messages

    def _execute_tool(self, tool_call: ToolCall) -> str:
        """Look up and execute the tool, returning the observation string."""
        tool = self.registry.get(tool_call.tool_name)
        if tool is None:
            return "Error: unknown tool '{}'. Available tools: {}".format(
                tool_call.tool_name,
                ", ".join(t.name for t in self.registry.list_tools()),
            )
        if tool.execute is None:
            return "Error: tool '{}' has no execute function.".format(tool_call.tool_name)

        # Coerce parameter types.
        kwargs: dict = {}
        param_types = {p.name: p.type for p in tool.params}
        for key, value in tool_call.arguments.items():
            expected = param_types.get(key, "str")
            kwargs[key] = coerce_param(value, expected)

        try:
            return tool.execute(**kwargs)
        except Exception as e:
            return "Error executing tool '{}': {}".format(tool_call.tool_name, e)