"""Prompt building utilities.

This module assembles:
- A deterministic system prompt (internal instructions + user system prompt +
  optional tool descriptions)
- Provider-agnostic message history plus the current user query
"""

from __future__ import annotations

import json
import logging
from typing import Any

from mcpbridge.config import PromptConfig
from mcpbridge.registry import MCPTool

logger = logging.getLogger("mcpbridge")


class PromptBuilder:
    """Build system prompts and messages for an LLM run."""

    INTERNAL_INSTRUCTIONS = """You have access to external tools via an MCP server connection.
Rules for tool use:
- Only call tools when needed to answer the user's request
- Always pass arguments that exactly match the tool's input schema
- If a tool returns an error, report it clearly; do not fabricate results
- You may call multiple tools in sequence or parallel as needed
- When you have enough information, stop calling tools and give your final answer
"""

    def __init__(self, prompt_config: PromptConfig) -> None:
        self._cfg = prompt_config
        self._context: dict[str, Any] = dict(prompt_config.context_vars)
        self._tools: list[MCPTool] = []

        # If True, PromptBuilder injects the system prompt into messages[0] and
        # returns system_prompt=None.
        self._inject_system_into_first_message = False

    def set_system_in_messages(self, inject: bool) -> None:
        """Configure whether system prompt is injected into messages[0]."""
        self._inject_system_into_first_message = bool(inject)

    def set_tools_context(self, tools: list[MCPTool]) -> None:
        """Set the tool context used for tool-description injection."""
        self._tools = list(tools)

    def set_context(self, **kwargs: Any) -> None:
        """Set runtime variables for `{placeholder}` interpolation."""
        self._context.update(kwargs)

    def build_system_prompt(self) -> str:
        """Assemble the system prompt in the required order."""
        parts: list[str] = []

        if self._cfg.inject_internal_instructions:
            parts.append(self.INTERNAL_INSTRUCTIONS.strip())

        sys_prompt = self._cfg.system or ""
        if self._cfg.interpolate and sys_prompt:
            try:
                sys_prompt = sys_prompt.format(**self._context)
            except KeyError as e:
                raise KeyError(f"Missing context var for system prompt: {e}") from e
        if sys_prompt:
            parts.append(sys_prompt.strip())

        if self._cfg.inject_tool_descriptions and self._tools:
            tool_lines: list[str] = []
            for t in self._tools:
                schema_snippet = json.dumps(t.input_schema, ensure_ascii=False)
                tool_lines.append(
                    f"Tool: {t.namespaced_name}\nDescription: {t.description}\nInput schema: {schema_snippet}"
                )
            parts.append("\n\n".join(tool_lines))

        return "\n\n".join(parts).strip()

    def build_messages(
        self,
        history: list[dict[str, Any]],
        user_query: str,
        user_prefix: str = "",
        user_suffix: str = "",
    ) -> tuple[str | None, list[dict[str, Any]]]:
        """Create system_prompt (or None) and an updated messages list."""
        history_copy = list(history or [])

        full_user_query = f"{user_prefix}{user_query}{user_suffix}"
        if self._cfg.user_prefix:
            full_user_query = f"{self._cfg.user_prefix}{full_user_query}"
        if self._cfg.user_suffix:
            full_user_query = f"{full_user_query}{self._cfg.user_suffix}"

        # Append user message.
        history_copy.append({"role": "user", "content": full_user_query})

        sys_prompt = self.build_system_prompt()
        if self._inject_system_into_first_message:
            if history_copy:
                history_copy[0] = {"role": "system", "content": sys_prompt}
            else:
                history_copy = [{"role": "system", "content": sys_prompt}, {"role": "user", "content": full_user_query}]
            return None, history_copy

        return (sys_prompt if sys_prompt else None), history_copy

