"""
Reasoning nodes for advanced agentic behaviors (ReAct, Planning).
"""

from typing import Any, Dict, List, Optional
import json
import re
from ybagent.node import BaseNode, NodeType, NodeInput, NodeOutput
from ybagent.nodes.agent_nodes import CodeAgentNode
from ybagent.utils.display import DisplayManager

class ReActAgentNode(BaseNode):
    """
    ReAct (Reasoning + Acting) Agent Node.
    
    This node implements a reasoning loop:
    1. Thought: Analyze the task and current state.
    2. Action: Decide on an action (tool call) or final answer.
    3. Observation: Execute the action and observe the result.
    4. Repeat until task is solved or max steps reached.
    """
    
    def __init__(
        self,
        id: str,
        tools: List[BaseNode] = None,
        max_steps: int = 10,
        llm_model: str = "gpt-4",
        llm_api_key: str = "",
        provider: str = "openai",
        **config
    ):
        """
        Initialize ReAct Agent.
        
        Args:
            id: Node ID
            tools: List of nodes that can be used as tools
            max_steps: Maximum reasoning steps
            llm_model: LLM model
            llm_api_key: API key
            provider: LLM provider
            **config: Additional config
        """
        config.update({
            "max_steps": max_steps,
            "llm_model": llm_model,
            "llm_api_key": llm_api_key,
            "provider": provider
        })
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.REACT_AGENT, **config)
        
        self.tools = {tool.metadata.id: tool for tool in (tools or [])}
        
        # Reuse CodeAgentNode's generation logic for convenience (or we could implement a separate one)
        # We'll create a helper CodeAgentNode internally just for LLM calls if needed, 
        # but here we'll implement a simple LLM caller directly to keep it self-contained.
        self.code_agent = CodeAgentNode(
            id=f"{id}_internal",
            sandbox_type="restricted", # Lightweight sandbox for internal logic if needed
            provider=provider,
            llm_model=llm_model,
            llm_api_key=llm_api_key
        )

    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="task", type="str", required=True, description="Goal/Task"),
            NodeInput(name="context", type="dict", required=False, description="Context")
        ]
        self.metadata.outputs = [
            NodeOutput(name="result", type="any", description="Final Answer"),
            NodeOutput(name="history", type="list", description="Reasoning Trace")
        ]

    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        task = inputs.get("task")
        history = []
        
        DisplayManager.header(f"🧠 ReAct Agent: {self.metadata.name}", level=3)
        
        # Display task cleanly (handle both string and dict inputs)
        if isinstance(task, dict):
            task_text = task.get('task', str(task))
        else:
            task_text = str(task)
        DisplayManager.info(f"Task: {task_text}")
        
        
        for step in range(self.metadata.config["max_steps"]):
            DisplayManager.step_start(step + 1, "Reasoning & Action")
            
            # 1. Generate Thought and Action
            prompt = self._build_prompt(task, history)
            response = self._query_llm(prompt)
            
            # Parse response
            thought, action, action_input = self._parse_response(response)
            
            step_record = {
                "step": step + 1,
                "thought": thought,
                "action": action,
                "action_input": action_input,
                "response": response
            }
            
            DisplayManager.thought(thought)
            
            if action == "Final Answer":
                DisplayManager.success(f"Final Answer: {action_input}")
                step_record["observation"] = "Task Completed"
                history.append(step_record)
                return {
                    "result": action_input,
                    "history": history
                }
            
            # 2. Execute Action
            DisplayManager.action(action, action_input)
            observation = self._execute_tool(action, action_input, context)
            
            # Format observation for better display
            formatted_observation = self._format_observation(observation)
            DisplayManager.observation(formatted_observation)
            
            step_record["observation"] = observation  # Store raw for history
            history.append(step_record)
            
        DisplayManager.error("Max steps reached without final answer.")
        return {
            "result": "Max steps reached without final answer.",
            "history": history
        }

    def _build_prompt(self, task: str, history: List[Dict]) -> str:
        tool_desc = "\n".join([f"- {name}: {tool.metadata.description}" for name, tool in self.tools.items()])
        
        history_str = ""
        for item in history:
            history_str += f"Thought: {item['thought']}\nAction: {item['action']}\nAction Input: {item['action_input']}\nObservation: {item['observation']}\n\n"
            
        return f"""Answer the following questions as best you can. You have access to the following tools:

{tool_desc}

Use the following format:

Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{', '.join(self.tools.keys())}] or "Final Answer"
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Action: Final Answer
Action Input: the final answer to the original input question

Begin!

Question: {task}
{history_str}Thought:"""

    def _parse_response(self, response: str):
        """Parse LLM response to extract thought, action, and action input."""
        # Extract thought (everything before "Action:")
        thought_match = re.search(r"(.*?)(?:\nAction:|\Z)", response, re.DOTALL)
        thought = thought_match.group(1).strip() if thought_match else response.strip()
        
        # Extract action
        action_match = re.search(r"Action:\s*(.*?)(?:\n|$)", response)
        action = action_match.group(1).strip() if action_match else "Final Answer"
        
        # Extract action input - try multiple patterns
        action_input = ""
        
        # Pattern 1: "Action Input: <value>"
        input_match = re.search(r"Action Input:\s*(.*?)(?:\n\n|\Z)", response, re.DOTALL)
        if input_match:
            action_input = input_match.group(1).strip()
        # Pattern 2: If Action is "Final Answer", try to get the answer from the thought
        elif action == "Final Answer":
            # Look for the answer in the thought itself
            # Common patterns: "the answer is X", "Final Answer: X", or just use the last line
            answer_patterns = [
                r"(?:answer is|result is|equals?)\s*:?\s*(\d+\.?\d*)",
                r"Final Answer:\s*(.+)",
                r"(\d+\.?\d*)\s*$"  # Just a number at the end
            ]
            for pattern in answer_patterns:
                match = re.search(pattern, thought, re.IGNORECASE)
                if match:
                    action_input = match.group(1).strip()
                    break
        
        return thought, action, action_input

    def _execute_tool(self, tool_name: str, tool_input: str, context: Dict) -> Any:
        if tool_name not in self.tools:
            return f"Error: Tool '{tool_name}' not found."
        
        tool = self.tools[tool_name]
        try:
            # Prepare input for the tool
            # This is a simplification; real mapping might be more complex
            tool_inputs = {"input": tool_input, "task": tool_input} 
            return tool.execute(tool_inputs, context)
        except Exception as e:
            return f"Error executing tool: {e}"
    
    def _format_observation(self, observation: Any) -> str:
        """Format observation for display - convert dicts/objects to pretty JSON."""
        if isinstance(observation, dict):
            # Extract the most important information
            if 'result' in observation:
                result = observation['result']
                logs = observation.get('logs', {})
                
                # Build a clean, structured output
                parts = []
                parts.append(f"✓ Result: {result}")
                
                if logs.get('success'):
                    parts.append(f"⏱️  Execution Time: {logs.get('execution_time', 0):.4f}s")
                    if logs.get('repaired'):
                        parts.append(f"🔧 Self-Corrected: Yes ({logs.get('attempts', 1)} attempts)")
                
                if 'code' in observation:
                    parts.append(f"\n📝 Code:\n{observation['code']}")
                
                return "\n".join(parts)
            else:
                # Generic dict formatting
                return json.dumps(observation, indent=2, ensure_ascii=False)
        elif isinstance(observation, (list, tuple)):
            return json.dumps(observation, indent=2, ensure_ascii=False)
        else:
            return str(observation)

    def _query_llm(self, prompt: str) -> str:
        # Reuse CodeAgentNode's generation method logic or call it directly
        # For now, we'll use a simplified version of what CodeAgentNode does
        # or actually just use the code_agent instance we created to generate "code" which is actually text here
        
        # Hack: We use the _generate_code method but treat the prompt as the task
        # We need to bypass the "You are a python code generator" system prompt though.
        # So we will implement a direct call here similar to CodeAgentNode but with different prompts.
        
        # For brevity in this implementation, let's assume we can use the internal code agent 
        # but we need to override the system prompt. 
        # Since CodeAgentNode doesn't easily allow system prompt override without inheritance change,
        # let's just copy the httpx logic here for now.
        
        provider = self.metadata.config.get("provider", "openai")
        model = self.metadata.config.get("llm_model", "gpt-4")
        api_key = self.metadata.config.get("llm_api_key", "")
        
        import httpx
        
        if provider == "openai":
            response = httpx.post(
                "https://api.openai.com/v1/chat/completions",
                headers={"Authorization": f"Bearer {api_key}"},
                json={
                    "model": model,
                    "messages": [{"role": "user", "content": prompt}],
                    "temperature": 0.0
                },
                timeout=30.0
            )
            return response.json()["choices"][0]["message"]["content"]
        
        elif provider == "ollama":
            base_url = "http://localhost:11434"
            try:
                response = httpx.post(
                    f"{base_url}/api/generate",
                    json={
                        "model": model,
                        "prompt": prompt,
                        "stream": False
                    },
                    timeout=60.0
                )
                response.raise_for_status()
                result = response.json()
                if "response" not in result:
                    raise ValueError(f"Ollama response missing 'response' key. Got: {result}")
                return result["response"]
            except httpx.HTTPStatusError as e:
                if e.response.status_code == 404:
                    raise ValueError(f"Ollama model '{model}' not found. Please run: ollama pull {model}")
                raise ValueError(f"Ollama API error: {e.response.status_code} - {e.response.text}")
            except httpx.RequestError as e:
                raise ValueError(f"Failed to connect to Ollama at {base_url}: {e}")
            
        return "Error: Provider not supported"
