"""
Agentic RAG Implementation

Implements RAG with autonomous agent capabilities:
- Tool use for dynamic information retrieval
- Multi-step reasoning with planning
- Self-directed exploration
- Action-observation loops
- Dynamic strategy selection

Based on ReAct, Toolformer, and agent research from 2024.
"""

import json
import logging
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional, Union

from ...workflow.builder import WorkflowBuilder
from ..ai.llm_agent import LLMAgentNode
from ..api.rest import RESTClientNode
from ..base import Node, NodeParameter, register_node
from ..code.python import PythonCodeNode
from ..data.sql import SQLDatabaseNode
from ..logic.workflow import WorkflowNode

logger = logging.getLogger(__name__)


@register_node()
class AgenticRAGNode(WorkflowNode):
    """
    Agentic RAG with Tool Use and Reasoning

    Implements autonomous agent capabilities for complex RAG tasks requiring
    multiple steps, external tools, and dynamic reasoning.

    When to use:
    - Best for: Complex research tasks, multi-source queries, dynamic exploration
    - Not ideal for: Simple lookups, static document sets
    - Performance: 3-10 seconds depending on reasoning steps
    - Quality improvement: 50-80% for complex analytical tasks

    Key features:
    - ReAct-style reasoning (Thought-Action-Observation loops)
    - Dynamic tool selection and use
    - Multi-step planning and execution
    - Self-directed information gathering
    - Verification and fact-checking

    Example:
        agentic_rag = AgenticRAGNode(
            tools=["search", "calculator", "database", "code_executor"],
            max_reasoning_steps=5
        )

        # Query: "Compare the revenue growth of tech companies in 2023 vs 2022"
        # Agent will:
        # 1. Plan the research approach
        # 2. Search for financial data
        # 3. Query databases for specific numbers
        # 4. Use calculator for growth calculations
        # 5. Synthesize findings with citations

        result = await agentic_rag.run(
            documents=financial_docs,
            query="Compare the revenue growth of tech companies in 2023 vs 2022"
        )

    Parameters:
        tools: List of available tools (search, api, database, etc.)
        max_reasoning_steps: Maximum reasoning iterations
        planning_strategy: How to plan actions (react, tree-of-thought)
        verification_enabled: Whether to verify findings

    Returns:
        answer: Final synthesized answer
        reasoning_trace: Complete thought-action-observation history
        tools_used: Which tools were utilized
        confidence: Agent's confidence in the answer
    """

    def __init__(
        self,
        name: str = "agentic_rag",
        tools: List[str] = None,
        max_reasoning_steps: int = 5,
        planning_strategy: str = "react",
        verification_enabled: bool = True,
    ):
        self.tools = tools or ["search", "calculator", "database"]
        self.max_reasoning_steps = max_reasoning_steps
        self.planning_strategy = planning_strategy
        self.verification_enabled = verification_enabled
        super().__init__(name, self._create_workflow())

    def _create_workflow(self) -> WorkflowNode:
        """Create agentic RAG workflow"""
        builder = WorkflowBuilder()

        # Planning agent
        planner_id = builder.add_node(
            "LLMAgentNode",
            node_id="planner_agent",
            config={
                "system_prompt": f"""You are a research planning agent. Given a query, create a step-by-step plan.

Available tools: {', '.join(self.tools)}

For each step, specify:
1. What information is needed
2. Which tool to use
3. Expected outcome

Return JSON:
{{
    "plan": [
        {{"step": 1, "action": "search", "query": "...", "purpose": "..."}},
        {{"step": 2, "action": "calculate", "expression": "...", "purpose": "..."}}
    ],
    "complexity": "simple|moderate|complex",
    "estimated_steps": 3
}}""",
                "model": "gpt-4",
            },
        )

        # ReAct reasoning loop
        react_agent_id = builder.add_node(
            "LLMAgentNode",
            node_id="react_agent",
            config={
                "system_prompt": f"""You are a ReAct agent that reasons step-by-step and uses tools.

Available tools:
- search(query): Search documents or web
- calculate(expression): Perform calculations
- database(query): Query structured data
- verify(claim): Fact-check a claim

Format your response:
Thought: [reasoning about what to do next]
Action: [tool_name(parameters)]
Observation: [I'll fill this in]

Continue until you have enough information to answer.
End with:
Answer: [final comprehensive answer]

Maximum steps: {self.max_reasoning_steps}""",
                "model": "gpt-4",
            },
        )

        # Tool executor
        tool_executor_id = builder.add_node(
            "PythonCodeNode",
            node_id="tool_executor",
            config={
                "code": """
import re
import json
from datetime import datetime

def execute_tool(action_string, documents, context):
    '''Execute tool based on action string'''
    # Parse action
    match = re.match(r'(\w+)\((.*)\)', action_string.strip())
    if not match:
        return {"error": "Invalid action format"}

    tool_name = match.group(1)
    params = match.group(2).strip('"\'')

    results = {}

    if tool_name == "search":
        # Search through documents
        query_words = set(params.lower().split())
        search_results = []

        for doc in documents[:50]:  # Limit for performance
            content = doc.get("content", "").lower()
            title = doc.get("title", "").lower()

            # Score based on word overlap
            doc_words = set(content.split())
            title_words = set(title.split())

            content_score = len(query_words & doc_words) / len(query_words) if query_words else 0
            title_score = len(query_words & title_words) / len(query_words) if query_words else 0

            total_score = content_score + (title_score * 2)  # Title matches weighted higher

            if total_score > 0:
                search_results.append({
                    "title": doc.get("title", "Untitled"),
                    "excerpt": content[:200] + "...",
                    "score": total_score,
                    "id": doc.get("id", "unknown")
                })

        # Sort by score
        search_results.sort(key=lambda x: x["score"], reverse=True)
        results = {
            "tool": "search",
            "query": params,
            "results": search_results[:5],
            "count": len(search_results)
        }

    elif tool_name == "calculate":
        # Safe calculation
        try:
            # Only allow basic math operations
            safe_dict = {"__builtins__": None}
            safe_dict.update({
                "abs": abs, "round": round, "min": min, "max": max,
                "sum": sum, "len": len, "pow": pow
            })

            result = eval(params, safe_dict)
            results = {
                "tool": "calculate",
                "expression": params,
                "result": result
            }
        except Exception as e:
            results = {
                "tool": "calculate",
                "error": str(e)
            }

    elif tool_name == "database":
        # Simulated database query
        if "revenue" in params.lower():
            # Mock financial data
            results = {
                "tool": "database",
                "query": params,
                "results": [
                    {"company": "TechCorp", "revenue_2022": 100, "revenue_2023": 120},
                    {"company": "DataInc", "revenue_2022": 80, "revenue_2023": 95},
                    {"company": "CloudCo", "revenue_2022": 60, "revenue_2023": 85}
                ]
            }
        else:
            results = {
                "tool": "database",
                "query": params,
                "results": []
            }

    elif tool_name == "verify":
        # Fact verification (simplified)
        confidence = 0.85 if "true" not in params.lower() else 0.95
        results = {
            "tool": "verify",
            "claim": params,
            "verified": confidence > 0.8,
            "confidence": confidence,
            "sources": ["Document analysis", "Cross-reference check"]
        }

    else:
        results = {"error": f"Unknown tool: {tool_name}"}

    return results

# Execute current action
reasoning_state = reasoning_state
documents = documents

current_action = reasoning_state.get("current_action", "")
if current_action:
    observation = execute_tool(current_action, documents, reasoning_state)
else:
    observation = {"error": "No action specified"}

result = {
    "tool_result": observation,
    "timestamp": datetime.now().isoformat()
}
"""
            },
        )

        # Reasoning state manager
        state_manager_id = builder.add_node(
            "PythonCodeNode",
            node_id="state_manager",
            config={
                "code": f"""
# Manage reasoning state across iterations
import json

def update_reasoning_state(state, new_response, tool_result=None):
    '''Update state with new reasoning step'''
    if not state:
        state = {{
            "steps": [],
            "current_step": 0,
            "completed": False,
            "final_answer": None
        }}

    # Parse response for thought/action/answer
    lines = new_response.strip().split('\\n')
    current_thought = None
    current_action = None
    final_answer = None

    for line in lines:
        if line.startswith("Thought:"):
            current_thought = line[8:].strip()
        elif line.startswith("Action:"):
            current_action = line[7:].strip()
        elif line.startswith("Answer:"):
            final_answer = line[7:].strip()
            state["completed"] = True

    # Add step
    step = {{
        "step_number": state["current_step"] + 1,
        "thought": current_thought,
        "action": current_action,
        "observation": tool_result.get("tool_result") if tool_result else None
    }}

    state["steps"].append(step)
    state["current_step"] += 1
    state["current_action"] = current_action
    state["final_answer"] = final_answer

    # Check if we've reached max steps
    if state["current_step"] >= {self.max_reasoning_steps}:
        state["completed"] = True
        if not state["final_answer"]:
            state["final_answer"] = "Reached maximum reasoning steps. Based on gathered information..."

    return state

# Process current iteration
plan = plan.get("response") if isinstance(plan, dict) else plan
reasoning_response = reasoning_response.get("response") if isinstance(reasoning_response, dict) else reasoning_response
tool_result = tool_result if "tool_result" in locals() else None

# Initialize or update state
if "reasoning_state" not in locals() or not reasoning_state:
    reasoning_state = None

reasoning_state = update_reasoning_state(
    reasoning_state,
    reasoning_response,
    tool_result
)

# Prepare context for next iteration
context_for_agent = ""
for step in reasoning_state["steps"]:
    if step["thought"]:
        context_for_agent += f"\\nThought: {{step['thought']}}"
    if step["action"]:
        context_for_agent += f"\\nAction: {{step['action']}}"
    if step["observation"]:
        context_for_agent += f"\\nObservation: {{step['observation']}}"

result = {{
    "reasoning_state": reasoning_state,
    "context_for_agent": context_for_agent,
    "continue_reasoning": not reasoning_state["completed"]
}}
"""
            },
        )

        # Verification agent (if enabled)
        if self.verification_enabled:
            verifier_id = builder.add_node(
                "LLMAgentNode",
                node_id="verifier_agent",
                config={
                    "system_prompt": """You are a fact-checking agent. Verify the accuracy of the answer.

Check for:
1. Factual accuracy
2. Logical consistency
3. Completeness
4. Source reliability

Return JSON:
{
    "verified": true/false,
    "confidence": 0.0-1.0,
    "issues": ["list of any issues found"],
    "suggestions": ["improvements if needed"]
}""",
                    "model": "gpt-4",
                },
            )

        # Result synthesizer
        synthesizer_id = builder.add_node(
            "PythonCodeNode",
            node_id="result_synthesizer",
            config={
                "code": """
# Synthesize final results
reasoning_state = reasoning_state
verification = verification if "verification" in locals() else None
query = query

# Extract tool usage
tools_used = []
for step in reasoning_state["steps"]:
    if step["observation"] and "tool" in step["observation"]:
        tools_used.append(step["observation"]["tool"])

# Calculate confidence
base_confidence = 0.7
confidence_boost = min(0.3, len(tools_used) * 0.1)
if verification and verification.get("response"):
    verification_data = verification["response"]
    if isinstance(verification_data, str):
        import json
        try:
            verification_data = json.loads(verification_data)
        except:
            verification_data = {"confidence": 0.8}

    verification_confidence = verification_data.get("confidence", 0.8)
    final_confidence = (base_confidence + confidence_boost) * verification_confidence
else:
    final_confidence = base_confidence + confidence_boost

# Build reasoning trace
reasoning_trace = []
for step in reasoning_state["steps"]:
    trace_entry = {
        "step": step["step_number"],
        "thought": step["thought"],
        "action": step["action"],
        "observation": step["observation"]
    }
    reasoning_trace.append(trace_entry)

result = {
    "agentic_rag_result": {
        "query": query,
        "answer": reasoning_state["final_answer"],
        "reasoning_trace": reasoning_trace,
        "tools_used": list(set(tools_used)),
        "confidence": final_confidence,
        "total_steps": len(reasoning_state["steps"]),
        "verification": verification.get("response") if verification else None,
        "metadata": {
            "planning_strategy": "{self.planning_strategy}",
            "max_steps": {self.max_reasoning_steps},
            "completed_successfully": reasoning_state["completed"]
        }
    }
}
"""
            },
        )

        # Connect workflow
        # Planning phase
        builder.add_connection(planner_id, "response", state_manager_id, "plan")

        # ReAct loop connections
        builder.add_connection(
            state_manager_id, "context_for_agent", react_agent_id, "additional_context"
        )
        builder.add_connection(
            react_agent_id, "response", state_manager_id, "reasoning_response"
        )
        builder.add_connection(
            state_manager_id, "reasoning_state", tool_executor_id, "reasoning_state"
        )
        builder.add_connection(
            tool_executor_id, "tool_result", state_manager_id, "tool_result"
        )

        # Loop control - continue if not completed
        builder.add_connection(
            state_manager_id, "continue_reasoning", react_agent_id, "_continue_if_true"
        )

        # Verification (if enabled)
        if self.verification_enabled:
            builder.add_connection(
                state_manager_id, "reasoning_state", verifier_id, "answer_to_verify"
            )
            builder.add_connection(
                verifier_id, "response", synthesizer_id, "verification"
            )

        # Final synthesis
        builder.add_connection(
            state_manager_id, "reasoning_state", synthesizer_id, "reasoning_state"
        )

        return builder.build(name="agentic_rag_workflow")


@register_node()
class ToolAugmentedRAGNode(Node):
    """
    Tool-Augmented RAG Node

    Enhances RAG with specific tool capabilities for specialized tasks.

    When to use:
    - Best for: Domain-specific queries requiring specialized tools
    - Not ideal for: General knowledge questions
    - Performance: 1-5 seconds depending on tools used
    - Accuracy: High for tool-supported domains

    Example:
        tool_rag = ToolAugmentedRAGNode(
            tool_registry={
                "calculator": calculate_func,
                "unit_converter": convert_units,
                "date_calculator": date_math
            }
        )

    Parameters:
        tool_registry: Dict of tool_name -> callable
        auto_detect_tools: Automatically detect needed tools
        fallback_strategy: What to do if tools fail

    Returns:
        answer: Tool-augmented response
        tools_invoked: List of tools used
        tool_outputs: Results from each tool
    """

    def __init__(
        self,
        name: str = "tool_augmented_rag",
        tool_registry: Dict[str, Callable] = None,
        auto_detect_tools: bool = True,
    ):
        self.tool_registry = tool_registry or {}
        self.auto_detect_tools = auto_detect_tools
        super().__init__(name)

    def get_parameters(self) -> Dict[str, NodeParameter]:
        return {
            "query": NodeParameter(
                name="query",
                type=str,
                required=True,
                description="Query requiring tool augmentation",
            ),
            "documents": NodeParameter(
                name="documents",
                type=list,
                required=False,
                description="Reference documents",
            ),
            "context": NodeParameter(
                name="context",
                type=dict,
                required=False,
                description="Additional context for tools",
            ),
        }

    def run(self, **kwargs) -> Dict[str, Any]:
        """Execute tool-augmented RAG"""
        query = kwargs.get("query", "")
        documents = kwargs.get("documents", [])
        context = kwargs.get("context", {})

        # Detect required tools
        required_tools = self._detect_required_tools(query)

        # Execute tools
        tool_outputs = {}
        for tool_name in required_tools:
            if tool_name in self.tool_registry:
                try:
                    tool_func = self.tool_registry[tool_name]
                    tool_outputs[tool_name] = tool_func(query, context)
                except Exception as e:
                    logger.error(f"Tool {tool_name} failed: {e}")
                    tool_outputs[tool_name] = {"error": str(e)}

        # Augment response with tool outputs
        augmented_answer = self._synthesize_with_tools(query, documents, tool_outputs)

        return {
            "answer": augmented_answer,
            "tools_invoked": list(required_tools),
            "tool_outputs": tool_outputs,
            "confidence": 0.9 if tool_outputs else 0.7,
        }

    def _detect_required_tools(self, query: str) -> List[str]:
        """Detect which tools are needed for the query"""
        required = []

        query_lower = query.lower()

        # Simple keyword detection (would use NER/classification in production)
        if any(
            word in query_lower for word in ["calculate", "compute", "sum", "average"]
        ):
            required.append("calculator")

        if any(word in query_lower for word in ["convert", "unit", "measurement"]):
            required.append("unit_converter")

        if any(word in query_lower for word in ["date", "days", "weeks", "months"]):
            required.append("date_calculator")

        return required

    def _synthesize_with_tools(
        self, query: str, documents: List[Dict], tool_outputs: Dict[str, Any]
    ) -> str:
        """Synthesize answer using tool outputs"""
        # In production, would use LLM for synthesis
        answer_parts = [f"Based on analysis of {len(documents)} documents"]

        if tool_outputs:
            answer_parts.append("and computational tools:")
            for tool, output in tool_outputs.items():
                if "error" not in output:
                    answer_parts.append(f"\n- {tool}: {output}")

        answer_parts.append(
            f"\nThe answer to '{query}' has been computed with tool assistance."
        )

        return " ".join(answer_parts)


@register_node()
class ReasoningRAGNode(WorkflowNode):
    """
    Multi-Step Reasoning RAG

    Implements complex reasoning chains for analytical queries.

    When to use:
    - Best for: Complex analytical questions, multi-step problems
    - Not ideal for: Simple factual queries
    - Performance: 2-8 seconds depending on reasoning depth
    - Quality: Superior for questions requiring logic and analysis

    Example:
        reasoning_rag = ReasoningRAGNode(
            reasoning_depth=3,
            strategy="chain_of_thought"
        )

        # Query: "If Company A grows 20% annually and Company B grows 15%,
        #         when will A's revenue exceed B's if B starts 50% larger?"
        # Will break down into steps and reason through the math

    Parameters:
        reasoning_depth: Maximum reasoning steps
        strategy: Reasoning strategy (chain_of_thought, tree_of_thought)
        verify_logic: Whether to verify logical consistency

    Returns:
        answer: Reasoned answer with steps
        reasoning_chain: Step-by-step logic
        assumptions: Assumptions made
        confidence: Confidence in reasoning
    """

    def __init__(
        self,
        name: str = "reasoning_rag",
        reasoning_depth: int = 3,
        strategy: str = "chain_of_thought",
    ):
        self.reasoning_depth = reasoning_depth
        self.strategy = strategy
        super().__init__(name, self._create_workflow())

    def _create_workflow(self) -> WorkflowNode:
        """Create reasoning RAG workflow"""
        builder = WorkflowBuilder()

        # Problem decomposer
        decomposer_id = builder.add_node(
            "LLMAgentNode",
            node_id="problem_decomposer",
            config={
                "system_prompt": f"""Break down complex problems into reasoning steps.

Strategy: {self.strategy}
Max depth: {self.reasoning_depth}

For each step specify:
1. What to determine
2. Required information
3. Logic/calculation needed

Return JSON:
{{
    "steps": [
        {{"step": 1, "goal": "...", "requires": ["..."], "approach": "..."}}
    ],
    "assumptions": ["list assumptions"],
    "complexity": "low|medium|high"
}}""",
                "model": "gpt-4",
            },
        )

        # Step-by-step reasoner
        step_reasoner_id = builder.add_node(
            "LLMAgentNode",
            node_id="step_reasoner",
            config={
                "system_prompt": """Execute one reasoning step at a time.

Given:
- Current step goal
- Available information
- Previous steps' results

Provide:
- Logical reasoning
- Calculations if needed
- Conclusion for this step
- What's needed next

Be explicit about your logic.""",
                "model": "gpt-4",
            },
        )

        # Logic verifier
        verifier_id = builder.add_node(
            "LLMAgentNode",
            node_id="logic_verifier",
            config={
                "system_prompt": """Verify the logical consistency of reasoning.

Check:
1. Are all steps logically sound?
2. Do conclusions follow from premises?
3. Are calculations correct?
4. Are assumptions reasonable?

Rate confidence: 0.0-1.0""",
                "model": "gpt-4",
            },
        )

        # Connect workflow
        builder.add_connection(
            decomposer_id, "response", step_reasoner_id, "reasoning_plan"
        )

        # Multiple reasoning steps (simplified - would use loop in production)
        builder.add_connection(
            step_reasoner_id, "response", verifier_id, "reasoning_to_verify"
        )

        return builder.build(name="reasoning_rag_workflow")


# Export all agentic nodes
__all__ = ["AgenticRAGNode", "ToolAugmentedRAGNode", "ReasoningRAGNode"]
