"""
JARVIS — ReAct Agent
Implements the Thought → Action → Observation → Response loop.

Pattern:
  THOUGHT: Analyze the user's request, check what's needed
  ACTION:  Call a tool or the LLM directly
  OBSERVATION: Examine the result
  RESPONSE: Generate final answer, store to memory
"""

import re
import json
from typing import List, Dict, Optional, Any
from dataclasses import dataclass, field
from enum import Enum
from loguru import logger

from app.pipelines.llm_pipeline import (
    build_chat_prompt,
    generate_response,
    SYSTEM_PROMPT,
)
from app.tools.tool_executor import (
    execute_tool,
    get_tool_descriptions,
    TOOL_DEFINITIONS,
)
from app.memory.faiss_store import FAISSMemoryStore
from app.core.config import settings


# ─────────────────────────────────────────────
# Data structures
# ─────────────────────────────────────────────
class StepType(str, Enum):
    THOUGHT     = "thought"
    ACTION      = "action"
    OBSERVATION = "observation"
    RESPONSE    = "response"


@dataclass
class AgentStep:
    step_type: StepType
    content:   str
    tool_name: Optional[str] = None
    tool_arg:  Optional[str] = None
    metadata:  Dict = field(default_factory=dict)


@dataclass
class AgentResult:
    final_response: str
    steps:         List[AgentStep]
    used_tools:    List[str]
    memory_hits:   int
    iterations:    int


# ─────────────────────────────────────────────
# System prompt with ReAct instructions
# ─────────────────────────────────────────────
REACT_SYSTEM_PROMPT = f"""{SYSTEM_PROMPT}

{get_tool_descriptions()}

REASONING PROTOCOL:
When you need to use a tool, format your response EXACTLY like this:
  THOUGHT: [your reasoning about what to do]
  ACTION: [TOOL: tool_name] argument

When you have enough information to answer:
  THOUGHT: [your final reasoning]
  RESPONSE: [your complete answer to the user]

Rules:
- Always start with THOUGHT
- Use tools when you need current information, calculations, or code execution
- After an OBSERVATION, continue reasoning or give RESPONSE
- RESPONSE must be complete and helpful
- If no tool is needed, go directly to RESPONSE
"""


# ─────────────────────────────────────────────
# Tool call parser
# ─────────────────────────────────────────────
TOOL_PATTERN = re.compile(
    r"\[TOOL:\s*([a-zA-Z_]+)\]\s*(.+?)(?=\n|$)",
    re.IGNORECASE | re.DOTALL,
)
ACTION_PATTERN = re.compile(
    r"ACTION:\s*(.+?)(?=\nTHOUGHT|\nRESPONSE|\nOBSERVATION|$)",
    re.DOTALL,
)
RESPONSE_PATTERN = re.compile(
    r"RESPONSE:\s*(.+?)$",
    re.DOTALL,
)
THOUGHT_PATTERN = re.compile(
    r"THOUGHT:\s*(.+?)(?=\nACTION|\nRESPONSE|$)",
    re.DOTALL,
)


def _parse_tool_call(action_text: str) -> Optional[tuple]:
    """Parse tool name and argument from ACTION text."""
    match = TOOL_PATTERN.search(action_text)
    if match:
        return match.group(1).strip(), match.group(2).strip()
    return None


def _parse_response(text: str) -> Optional[str]:
    """Extract RESPONSE block from LLM output."""
    match = RESPONSE_PATTERN.search(text)
    if match:
        return match.group(1).strip()
    return None


# ─────────────────────────────────────────────
# ReAct Agent
# ─────────────────────────────────────────────
class ReActAgent:
    """
    Multi-step reasoning agent with tool use.
    Implements the ReAct (Reasoning + Acting) pattern.
    """

    def __init__(self):
        self.memory = FAISSMemoryStore.get_instance()

    def run(
        self,
        user_message: str,
        conversation_history: List[Dict[str, str]],
        session_id: Optional[str] = None,
        image_caption: Optional[str] = None,
        system_override: Optional[str] = None,
    ) -> AgentResult:
        """
        Run the full ReAct loop.

        Args:
            user_message: The user's input text
            conversation_history: Previous messages in this session
            session_id: Optional session ID for memory scoping
            image_caption: Optional image context string
            system_override: Optional system prompt override

        Returns:
            AgentResult with final response and trace
        """
        steps: List[AgentStep] = []
        used_tools: List[str] = []
        iterations = 0
        max_iter = settings.AGENT_MAX_ITERATIONS

        # ── Step 1: Retrieve relevant memories ────────────────────
        memory_context = self.memory.search_context(
            user_message,
            session_id=session_id,
        )
        memory_hits = len(self.memory.search(user_message, session_id=session_id))

        # ── Step 2: Check if tools are needed (lightweight) ────────
        needs_tool = self._classify_needs_tool(user_message)

        if not needs_tool:
            # Fast path: direct LLM response
            prompt = build_chat_prompt(
                user_message=user_message,
                history=conversation_history,
                memory_context=memory_context,
                image_caption=image_caption,
                system_override=system_override,
            )
            response_text = generate_response(prompt)

            steps.append(AgentStep(StepType.THOUGHT, "Direct response — no tools needed"))
            steps.append(AgentStep(StepType.RESPONSE, response_text))

            # Store interaction in memory
            self._store_interaction(user_message, response_text, session_id)

            return AgentResult(
                final_response=response_text,
                steps=steps,
                used_tools=[],
                memory_hits=memory_hits,
                iterations=1,
            )

        # ── Step 3: ReAct loop ─────────────────────────────────────
        # Build the agent prompt with ReAct instructions
        react_system = REACT_SYSTEM_PROMPT
        if memory_context:
            react_system += f"\n\n[MEMORY CONTEXT]\n{memory_context}"
        if image_caption:
            react_system += f"\n\n[IMAGE]\nCaption: {image_caption}"

        agent_history = list(conversation_history)
        current_message = user_message
        final_response = None

        # Accumulated context from tool observations
        observations_context = ""

        while iterations < max_iter and final_response is None:
            iterations += 1

            # Build prompt for this iteration
            if observations_context:
                augmented_message = (
                    f"{current_message}\n\n"
                    f"[PREVIOUS OBSERVATIONS]\n{observations_context}"
                )
            else:
                augmented_message = current_message

            prompt = build_chat_prompt(
                user_message=augmented_message,
                history=agent_history,
                system_override=react_system,
            )

            # Generate next step
            raw_output = generate_response(
                prompt,
                max_new_tokens=300,
                stop_strings=["OBSERVATION:", "[OBSERVATION]"],
            )

            logger.debug(f"Agent iteration {iterations}: {raw_output[:200]}")

            # Parse thought
            thought_match = THOUGHT_PATTERN.search(raw_output)
            thought_text = thought_match.group(1).strip() if thought_match else raw_output[:200]
            steps.append(AgentStep(StepType.THOUGHT, thought_text))

            # Check for RESPONSE
            response_text = _parse_response(raw_output)
            if response_text:
                final_response = response_text
                steps.append(AgentStep(StepType.RESPONSE, response_text))
                break

            # Check for ACTION / tool call
            action_match = ACTION_PATTERN.search(raw_output)
            if action_match:
                action_text = action_match.group(1).strip()
                tool_parsed = _parse_tool_call(action_text)

                if tool_parsed:
                    tool_name, tool_arg = tool_parsed
                    steps.append(AgentStep(
                        StepType.ACTION,
                        action_text,
                        tool_name=tool_name,
                        tool_arg=tool_arg,
                    ))

                    # Execute tool
                    tool_result = execute_tool(tool_name, tool_arg)
                    observation = tool_result.get("result", "No result")
                    used_tools.append(tool_name)

                    obs_step = AgentStep(
                        StepType.OBSERVATION,
                        observation,
                        tool_name=tool_name,
                        metadata={"success": tool_result.get("success", False)},
                    )
                    steps.append(obs_step)

                    # Add observation to context for next iteration
                    observations_context += f"\n[{tool_name} result]\n{observation}\n"

                    continue  # next iteration

            # No tool call and no response — extract whatever the LLM said
            # (graceful degradation)
            fallback = raw_output.replace("THOUGHT:", "").strip()
            if fallback:
                final_response = fallback
                steps.append(AgentStep(StepType.RESPONSE, fallback))
            else:
                final_response = "I processed your request but couldn't generate a clear response."
                steps.append(AgentStep(StepType.RESPONSE, final_response))

        # Max iterations reached without response
        if final_response is None:
            final_response = (
                "I've gathered the following information:\n" + observations_context
                if observations_context
                else "I encountered an issue processing your request. Please try again."
            )

        # Store interaction in memory
        self._store_interaction(user_message, final_response, session_id)

        return AgentResult(
            final_response=final_response,
            steps=steps,
            used_tools=used_tools,
            memory_hits=memory_hits,
            iterations=iterations,
        )

    def _classify_needs_tool(self, message: str) -> bool:
        """
        Lightweight heuristic to detect if tools are likely needed.
        Avoids the overhead of a full ReAct loop for simple queries.
        """
        message_lower = message.lower()

        tool_indicators = [
            # Web search
            "latest", "current", "today", "news", "weather", "price", "what is",
            "who is", "when did", "search", "find out", "look up",
            # Calculator
            "calculate", "compute", "what is", "how much", "solve", "math",
            "percent", "%", "=", "+", "-", "multiply", "divide",
            # Code
            "code", "python", "run", "execute", "program", "function", "write a script",
        ]

        # Check for explicit math expression patterns
        if re.search(r"\d+\s*[\+\-\*\/\^]\s*\d+", message):
            return True

        return any(indicator in message_lower for indicator in tool_indicators)

    def _store_interaction(
        self,
        user_message: str,
        assistant_response: str,
        session_id: Optional[str],
    ):
        """Store the interaction as a combined memory entry."""
        try:
            combined = f"User: {user_message}\nAssistant: {assistant_response}"
            self.memory.add(combined, session_id=session_id)
        except Exception as e:
            logger.warning(f"Failed to store interaction in memory: {e}")