"""MonitorOperativeAgent -- long-horizon agent with configurable strategies.

Extends ToolUsingAgent (not OperativeAgent) with four configurable strategy
axes for long-horizon benchmark evaluation:

1. **memory_extraction** -- how findings are persisted to memory
2. **observation_compression** -- how tool outputs are compressed
3. **retrieval_strategy** -- how prior context is recalled
4. **task_decomposition** -- how complex tasks are split

The agent also inherits cross-session state persistence from the
OperativeAgent pattern (session_store, memory_backend, operator_id).
"""

from __future__ import annotations

import json
import logging
import re
from typing import Any, List, Optional

from openjarvis.agents._stubs import AgentContext, AgentResult, ToolUsingAgent
from openjarvis.agents.prompt_loader import (
    load_few_shot_exemplars,
    load_system_prompt_override,
)
from openjarvis.core.events import EventBus
from openjarvis.core.registry import AgentRegistry
from openjarvis.core.types import Message, Role, ToolCall, ToolResult, _message_to_dict
from openjarvis.engine._stubs import InferenceEngine
from openjarvis.tools._stubs import BaseTool

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Valid strategy values
# ---------------------------------------------------------------------------

VALID_MEMORY_EXTRACTION = {"causality_graph", "scratchpad", "structured_json", "none"}
VALID_OBSERVATION_COMPRESSION = {"summarize", "truncate", "none"}
VALID_RETRIEVAL_STRATEGY = {"hybrid_with_self_eval", "keyword", "semantic", "none"}
VALID_TASK_DECOMPOSITION = {"phased", "monolithic", "hierarchical"}

# ---------------------------------------------------------------------------
# Default system prompt
# ---------------------------------------------------------------------------

MONITOR_OPERATIVE_SYSTEM_PROMPT = """\
You are a Monitor Operative Agent designed for long-horizon tasks.

## Capabilities
1. TOOLS: You have access to tools via native function calling. The list
   below shows what is available — invoke them through the function-calling
   API, not by writing tool names into your text response.
2. STATE: Your previous findings and state are automatically restored from memory.
3. MEMORY: Store important findings via memory_store; recall via memory_retrieve.

## Critical Operating Rule
Your training data is frozen and out of date. For ANY question about recent,
current, or evolving information, you MUST call a substantive retrieval tool
(web_search, memory_retrieve, or an equivalent) BEFORE composing a response.
Writing fact claims about recent events from memory alone produces
hallucinations and is a failure mode.

## Strategy
- Memory extraction: {memory_extraction}
- Observation compression: {observation_compression}
- Retrieval strategy: {retrieval_strategy}
- Task decomposition: {task_decomposition}

## Protocol
- Break complex tasks into phases and track progress
- Prefer substantive tools (web_search, memory_retrieve) over reasoning-only
  tools (think) — `think` does not gather new information, only reorganises
  what you already have
- Store causal relationships and key findings in memory
- Compress long tool outputs before adding to context
- Self-evaluate retrieved context for relevance
- Always persist state before finishing

{tool_descriptions}"""


@AgentRegistry.register("monitor_operative")
class MonitorOperativeAgent(ToolUsingAgent):
    """Long-horizon agent with configurable memory, compression, retrieval,
    and decomposition strategies.

    The four strategy axes control how the agent manages information across
    turns and sessions:

    - ``memory_extraction``: How findings are persisted (causality_graph,
      scratchpad, structured_json, none).
    - ``observation_compression``: How tool outputs are compressed before
      being added to context (summarize, truncate, none).
    - ``retrieval_strategy``: How prior context is recalled at the start
      of each run (hybrid_with_self_eval, keyword, semantic, none).
    - ``task_decomposition``: How complex tasks are broken down
      (phased, monolithic, hierarchical).
    """

    agent_id = "monitor_operative"
    accepts_tools = True
    _default_temperature = 0.3
    _default_max_tokens = 4096
    _default_max_turns = 25

    def __init__(
        self,
        engine: InferenceEngine,
        model: str,
        *,
        tools: Optional[List[BaseTool]] = None,
        bus: Optional[EventBus] = None,
        max_turns: Optional[int] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        system_prompt: Optional[str] = None,
        # Strategy parameters
        memory_extraction: str = "causality_graph",
        observation_compression: str = "summarize",
        retrieval_strategy: str = "hybrid_with_self_eval",
        task_decomposition: str = "phased",
        # State persistence (OperativeAgent pattern)
        operator_id: Optional[str] = None,
        session_store: Optional[Any] = None,
        memory_backend: Optional[Any] = None,
        interactive: bool = False,
        confirm_callback=None,
        **kwargs: Any,
    ) -> None:
        super().__init__(
            engine,
            model,
            tools=tools,
            bus=bus,
            max_turns=max_turns,
            temperature=temperature,
            max_tokens=max_tokens,
            interactive=interactive,
            confirm_callback=confirm_callback,
            prompt_builder=kwargs.get("prompt_builder"),
        )
        # Validate strategies
        if memory_extraction not in VALID_MEMORY_EXTRACTION:
            raise ValueError(
                f"Invalid memory_extraction {memory_extraction!r}, "
                f"must be one of {VALID_MEMORY_EXTRACTION}"
            )
        if observation_compression not in VALID_OBSERVATION_COMPRESSION:
            raise ValueError(
                f"Invalid observation_compression {observation_compression!r}, "
                f"must be one of {VALID_OBSERVATION_COMPRESSION}"
            )
        if retrieval_strategy not in VALID_RETRIEVAL_STRATEGY:
            raise ValueError(
                f"Invalid retrieval_strategy {retrieval_strategy!r}, "
                f"must be one of {VALID_RETRIEVAL_STRATEGY}"
            )
        if task_decomposition not in VALID_TASK_DECOMPOSITION:
            raise ValueError(
                f"Invalid task_decomposition {task_decomposition!r}, "
                f"must be one of {VALID_TASK_DECOMPOSITION}"
            )

        self._memory_extraction = memory_extraction
        self._observation_compression = observation_compression
        self._retrieval_strategy = retrieval_strategy
        self._task_decomposition = task_decomposition

        self._system_prompt = system_prompt
        self._operator_id = operator_id
        self._session_store = session_store
        self._memory_backend = memory_backend

    # ------------------------------------------------------------------
    # Main run loop
    # ------------------------------------------------------------------

    def run(
        self,
        input: str,
        context: Optional[AgentContext] = None,
        **kwargs: Any,
    ) -> AgentResult:
        """Execute the agent on *input* with the configured strategies."""
        self._emit_turn_start(input)

        # 1. Build system prompt with state context
        #    Priority: constructor arg > file override > hardcoded default
        sys_parts: list[str] = []
        if self._system_prompt:
            sys_parts.append(self._system_prompt)
        else:
            tool_desc = self._build_tool_descriptions()
            prompt_template = (
                load_system_prompt_override("monitor_operative")
                or MONITOR_OPERATIVE_SYSTEM_PROMPT
            )
            try:
                sys_parts.append(
                    prompt_template.format(
                        memory_extraction=self._memory_extraction,
                        observation_compression=self._observation_compression,
                        retrieval_strategy=self._retrieval_strategy,
                        task_decomposition=self._task_decomposition,
                        tool_descriptions=tool_desc,
                    ),
                )
            except KeyError:
                sys_parts.append(prompt_template)

        # 2. State recall from memory backend
        previous_state = self._recall_state()
        if previous_state:
            sys_parts.append(f"\n## Previous State\n{previous_state}")

        system_prompt = "\n\n".join(sys_parts) if sys_parts else None
        # Honor SOUL.md / MEMORY.md / USER.md persona files like `jarvis ask`,
        # appended so the monitor's own instructions are preserved (#376).
        system_prompt = self._apply_persona(system_prompt)

        # 3. Load session history
        session_messages = self._load_session()

        # 4. Build messages
        messages = self._build_operative_messages(
            input,
            context,
            system_prompt=system_prompt,
            session_messages=session_messages,
        )

        # 4b. Inject few-shot exemplars before the user input
        for ex in load_few_shot_exemplars("monitor_operative"):
            if ex.get("input") and ex.get("output"):
                messages.insert(-1, Message(role=Role.USER, content=ex["input"]))
                messages.insert(-1, Message(role=Role.ASSISTANT, content=ex["output"]))

        # 5. Run function-calling tool loop
        openai_tools = self._executor.get_openai_tools() if self._tools else []
        all_tool_results: list[ToolResult] = []
        turns = 0
        content = ""
        state_stored_by_tool = False
        total_usage: dict[str, int] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
        }

        for _turn in range(self._max_turns):
            turns += 1

            gen_kwargs: dict[str, Any] = {}
            if openai_tools:
                gen_kwargs["tools"] = openai_tools

            result = self._generate(messages, **gen_kwargs)
            usage = result.get("usage", {})
            for k in total_usage:
                total_usage[k] += usage.get(k, 0)
            content = result.get("content", "")
            # Strip think tags so they don't interfere with parsing
            content = self._strip_think_tags(content)
            raw_tool_calls = result.get("tool_calls", [])

            # --- Native function-calling path ---
            if raw_tool_calls:
                tool_calls = [
                    ToolCall(
                        id=tc.get("id", f"call_{i}"),
                        name=tc.get("name", ""),
                        arguments=tc.get("arguments", "{}"),
                    )
                    for i, tc in enumerate(raw_tool_calls)
                ]
                messages.append(
                    Message(
                        role=Role.ASSISTANT,
                        content=content,
                        tool_calls=tool_calls,
                    )
                )
            else:
                # --- Text-based fallback ---
                tool_info = self._extract_tool_call(content)
                if tool_info:
                    action, action_input = tool_info
                    messages.append(Message(role=Role.ASSISTANT, content=content))
                    tc = ToolCall(
                        id=f"text_call_{turns}",
                        name=action,
                        arguments=action_input,
                    )
                    tool_result = self._executor.execute(tc)
                    all_tool_results.append(tool_result)
                    observation_content = self._compress_observation(
                        tool_result.content
                    )
                    messages.append(
                        Message(
                            role=Role.USER,
                            content=f"Result: {observation_content}",
                        )
                    )
                    self._extract_and_store(tc.name, tool_result.content)
                    continue

                # No tool calls at all -> check continuation, then final answer
                content = self._check_continuation(result, messages)
                break

            # Execute each native tool call
            tool_calls_to_exec = tool_calls
            for tc in tool_calls_to_exec:
                # Loop guard check
                if self._loop_guard:
                    verdict = self._loop_guard.check_call(
                        tc.name,
                        tc.arguments,
                    )
                    if verdict.blocked:
                        tool_result = ToolResult(
                            tool_name=tc.name,
                            content=f"Loop guard: {verdict.reason}",
                            success=False,
                        )
                        all_tool_results.append(tool_result)
                        messages.append(
                            Message(
                                role=Role.TOOL,
                                content=tool_result.content,
                                tool_call_id=tc.id,
                                name=tc.name,
                            )
                        )
                        continue

                tool_result = self._executor.execute(tc)
                all_tool_results.append(tool_result)

                # Track explicit state storage
                if tc.name == "memory_store" and self._operator_id:
                    try:
                        args = json.loads(tc.arguments)
                        state_key = f"monitor_operative:{self._operator_id}:state"
                        if args.get("key", "") == state_key:
                            state_stored_by_tool = True
                    except (json.JSONDecodeError, TypeError) as exc:
                        logger.debug(
                            "Failed to parse tool call arguments"
                            " for state tracking: %s",
                            exc,
                        )

                # Compress observation if strategy requires it
                observation_content = self._compress_observation(tool_result.content)

                messages.append(
                    Message(
                        role=Role.TOOL,
                        content=observation_content,
                        tool_call_id=tc.id,
                        name=tc.name,
                    )
                )

                # Extract and store findings based on memory strategy
                self._extract_and_store(tc.name, tool_result.content)
        else:
            # Max turns exceeded
            self._save_session(input, content)
            msg_dicts = [_message_to_dict(m) for m in messages]
            return self._max_turns_result(
                all_tool_results,
                turns,
                content=content,
                metadata={**total_usage, "messages": msg_dicts},
            )

        # 6. Save session
        self._save_session(input, content)

        # 7. Auto-persist state if agent didn't do it explicitly
        if not state_stored_by_tool:
            self._auto_persist_state(content)

        self._emit_turn_end(turns=turns, content_length=len(content))
        msg_dicts = [_message_to_dict(m) for m in messages]
        return AgentResult(
            content=content,
            tool_results=all_tool_results,
            turns=turns,
            metadata={**total_usage, "messages": msg_dicts},
        )

    # ------------------------------------------------------------------
    # Text-based tool call extraction (fallback for non-function-calling models)
    # ------------------------------------------------------------------

    @staticmethod
    def _extract_tool_call(text: str) -> tuple[str, str] | None:
        """Extract tool call from text output.

        Supports three formats:
        1. Action: tool_name / Action Input: {"key": "value"}
        2. <tool_call>tool_name\\n$key=value</tool_call> (XML-style)
        3. <tool_name query="..."> or <tool_name>...</tool_name> (inline XML)
        """
        # Format 1: Action / Action Input
        action_match = re.search(r"Action:\s*(.+)", text, re.IGNORECASE)
        input_match = re.search(
            r"Action Input:\s*(.+?)(?=\n\n|\Z)", text, re.DOTALL | re.IGNORECASE
        )
        if action_match:
            return (
                action_match.group(1).strip(),
                input_match.group(1).strip() if input_match else "{}",
            )

        # Format 2: <tool_call>tool_name ... </tool_call>
        xml_match = re.search(
            r"<tool_call>\s*(\w+)\s*(.*?)</\w+>",
            text,
            re.DOTALL,
        )
        if xml_match:
            tool_name = xml_match.group(1).strip()
            raw_params = xml_match.group(2).strip()
            params: dict[str, Any] = {}
            for m in re.finditer(
                r"\$(\w+)=(.+?)(?=\$|\n<|</|$)", raw_params, re.DOTALL
            ):
                params[m.group(1)] = m.group(2).strip().rstrip("</>\n")
            for m in re.finditer(r"<(\w+)>(.*?)</\1>", raw_params, re.DOTALL):
                key, val = m.group(1), m.group(2).strip()
                try:
                    params[key] = int(val)
                except ValueError:
                    params[key] = val
            if not params:
                for m in re.finditer(
                    r"(\w+)\s*:\s*(.+?)(?=\n\w+\s*:|$)", raw_params, re.DOTALL
                ):
                    key, val = m.group(1), m.group(2).strip().strip("\"'")
                    try:
                        params[key] = int(val)
                    except ValueError:
                        params[key] = val
            if params:
                return (tool_name, json.dumps(params))
            return (tool_name, "{}")

        # Format 3: <web_search query="..."> or <tool_name>args</tool_name>
        # Handles Qwen-style XML tool output like <web_search query="...">
        inline_match = re.search(
            r"<(\w+)\s+(.*?)/?>",
            text,
            re.DOTALL,
        )
        if inline_match:
            tool_name = inline_match.group(1).strip()
            # Skip common non-tool tags
            if tool_name.lower() in ("think", "br", "hr", "p", "div", "span", "b", "i"):
                return None
            attr_str = inline_match.group(2).strip()
            params = {}
            for m in re.finditer(r'(\w+)=["\']([^"\']*)["\']', attr_str):
                params[m.group(1)] = m.group(2)
            # Also handle unquoted: <web_search query=something>
            if not params:
                for m in re.finditer(r"(\w+)=(\S+)", attr_str):
                    params[m.group(1)] = m.group(2).rstrip(">")
            if params:
                return (tool_name, json.dumps(params))
            return (tool_name, "{}")

        return None

    # ------------------------------------------------------------------
    # Message building
    # ------------------------------------------------------------------

    def _build_operative_messages(
        self,
        input: str,
        context: Optional[AgentContext],
        *,
        system_prompt: Optional[str] = None,
        session_messages: Optional[list[Message]] = None,
    ) -> list[Message]:
        """Build message list with system prompt, session history, and input."""
        messages: list[Message] = []
        if system_prompt:
            messages.append(Message(role=Role.SYSTEM, content=system_prompt))
        if session_messages:
            messages.extend(session_messages)
        if context and context.conversation.messages:
            messages.extend(context.conversation.messages)
        messages.append(Message(role=Role.USER, content=input))
        return messages

    def _build_tool_descriptions(self) -> str:
        """Build a text description of available tools for the system prompt."""
        if not self._tools:
            return ""
        from openjarvis.tools._stubs import build_tool_descriptions

        return build_tool_descriptions(self._tools)

    # ------------------------------------------------------------------
    # Strategy methods
    # ------------------------------------------------------------------

    def _compress_observation(self, content: str) -> str:
        """Compress a tool observation according to the compression strategy.

        - ``summarize``: If content exceeds 2000 chars, ask the LLM to
          summarize. Falls back to truncation if generation fails.
        - ``truncate``: Hard-truncate at 2000 chars with an ellipsis.
        - ``none``: Return content unchanged.
        """
        if self._observation_compression == "none":
            return content
        if self._observation_compression == "truncate":
            if len(content) > 2000:
                return content[:2000] + "\n... [truncated]"
            return content
        # "summarize"
        if len(content) <= 2000:
            return content
        try:
            summary_messages = [
                Message(
                    role=Role.SYSTEM,
                    content="Summarize the following tool output concisely, "
                    "preserving all key facts and data points.",
                ),
                Message(role=Role.USER, content=content[:8000]),
            ]
            result = self._generate(summary_messages)
            summary = result.get("content", "")
            if summary:
                return summary
        except Exception:
            logger.debug("Observation summarization failed, falling back to truncation")
        # Fallback to truncation
        return content[:2000] + "\n... [truncated]"

    def _extract_and_store(self, tool_name: str, content: str) -> None:
        """Extract findings from a tool result and store them.

        The extraction strategy depends on ``_memory_extraction``:

        - ``causality_graph``: Extract causal relationships via
          :meth:`_extract_causality` and store as KG triples.
        - ``scratchpad``: Append raw content to a scratchpad key in
          memory.
        - ``structured_json``: Attempt to parse JSON from the content
          and store structured data.
        - ``none``: Do nothing.
        """
        if self._memory_extraction == "none":
            return
        if not self._memory_backend:
            return

        if self._memory_extraction == "causality_graph":
            self._extract_causality(tool_name, content)
        elif self._memory_extraction == "scratchpad":
            self._store_scratchpad(tool_name, content)
        elif self._memory_extraction == "structured_json":
            self._store_structured(tool_name, content)

    def _extract_causality(self, tool_name: str, content: str) -> None:
        """Extract causal relationships from tool output and store them.

        Uses the LLM to identify cause-effect patterns, then stores
        them via the memory backend.
        """
        if not self._memory_backend or not content.strip():
            return
        # Only attempt extraction for substantial outputs
        if len(content) < 50:
            return
        try:
            extract_messages = [
                Message(
                    role=Role.SYSTEM,
                    content=(
                        "Extract causal relationships from the following tool "
                        "output. Return a JSON array of objects with 'cause', "
                        "'effect', and 'confidence' fields. If no causal "
                        "relationships are found, return an empty array []."
                    ),
                ),
                Message(role=Role.USER, content=content[:4000]),
            ]
            result = self._generate(extract_messages)
            raw = result.get("content", "")
            # Try to parse JSON from the response
            raw = raw.strip()
            if raw.startswith("```"):
                # Strip code fences
                lines = raw.split("\n")
                raw = "\n".join(lines[1:-1] if len(lines) > 2 else lines)
            relations = json.loads(raw)
            if isinstance(relations, list):
                operator_prefix = (
                    f"monitor_operative:{self._operator_id}"
                    if self._operator_id
                    else "monitor_operative"
                )
                for rel in relations[:10]:  # Cap at 10 per extraction
                    if isinstance(rel, dict) and "cause" in rel and "effect" in rel:
                        key = f"{operator_prefix}:causality:{rel['cause'][:50]}"
                        value = json.dumps(rel)
                        try:
                            self._memory_backend.store(key, value)
                        except Exception as exc:
                            logger.debug(
                                "Failed to store causality relation in memory: %s",
                                exc,
                            )
        except (json.JSONDecodeError, Exception):
            logger.debug(
                "Causality extraction failed for tool %s output",
                tool_name,
            )

    def _store_scratchpad(self, tool_name: str, content: str) -> None:
        """Append content to a scratchpad entry in memory."""
        if not self._memory_backend:
            return
        operator_prefix = (
            f"monitor_operative:{self._operator_id}"
            if self._operator_id
            else "monitor_operative"
        )
        key = f"{operator_prefix}:scratchpad:{tool_name}"
        # Truncate long content
        snippet = content[:1000] if len(content) > 1000 else content
        try:
            self._memory_backend.store(key, snippet)
        except Exception:
            logger.debug("Could not store scratchpad for tool %s", tool_name)

    def _store_structured(self, tool_name: str, content: str) -> None:
        """Try to parse JSON from tool output and store structured data."""
        if not self._memory_backend:
            return
        operator_prefix = (
            f"monitor_operative:{self._operator_id}"
            if self._operator_id
            else "monitor_operative"
        )
        try:
            data = json.loads(content)
            key = f"{operator_prefix}:structured:{tool_name}"
            self._memory_backend.store(key, json.dumps(data))
        except (json.JSONDecodeError, TypeError):
            # Not JSON -- store as plain text truncated
            key = f"{operator_prefix}:structured:{tool_name}"
            try:
                self._memory_backend.store(key, content[:1000])
            except Exception as exc:
                logger.debug(
                    "Failed to store structured data for tool %s: %s",
                    tool_name,
                    exc,
                )

    # ------------------------------------------------------------------
    # State persistence (OperativeAgent pattern)
    # ------------------------------------------------------------------

    def _recall_state(self) -> str:
        """Retrieve previous state from memory backend."""
        if not self._memory_backend or not self._operator_id:
            return ""
        state_key = f"monitor_operative:{self._operator_id}:state"
        try:
            result = self._memory_backend.retrieve(state_key)
            if result:
                return result if isinstance(result, str) else str(result)
        except Exception:
            logger.debug(
                "No previous state for monitor_operative %s",
                self._operator_id,
            )
        return ""

    def _load_session(self) -> list[Message]:
        """Load recent session history for this operator."""
        if not self._session_store or not self._operator_id:
            return []
        session_id = f"monitor_operative:{self._operator_id}"
        try:
            session = self._session_store.get_or_create(session_id)
            if hasattr(session, "messages") and session.messages:
                recent = session.messages[-10:]
                return [
                    Message(
                        role=Role(m.get("role", "user")),
                        content=m.get("content", ""),
                    )
                    for m in recent
                    if isinstance(m, dict)
                ]
        except Exception:
            logger.debug(
                "Could not load session for monitor_operative %s",
                self._operator_id,
            )
        return []

    def _save_session(self, input_text: str, response: str) -> None:
        """Save the tick's prompt and response to the session store."""
        if not self._session_store or not self._operator_id:
            return
        session_id = f"monitor_operative:{self._operator_id}"
        try:
            self._session_store.save_message(
                session_id,
                {"role": "user", "content": input_text},
            )
            self._session_store.save_message(
                session_id,
                {"role": "assistant", "content": response},
            )
        except Exception:
            logger.debug(
                "Could not save session for monitor_operative %s",
                self._operator_id,
            )

    def _auto_persist_state(self, content: str) -> None:
        """Auto-persist a state summary if agent didn't store explicitly."""
        if not self._memory_backend or not self._operator_id:
            return
        state_key = f"monitor_operative:{self._operator_id}:state"
        try:
            summary = content[:1000] if content else ""
            self._memory_backend.store(state_key, summary)
        except Exception:
            logger.debug(
                "Could not auto-persist state for monitor_operative %s",
                self._operator_id,
            )


__all__ = ["MonitorOperativeAgent", "MONITOR_OPERATIVE_SYSTEM_PROMPT"]
