"""
JARVIS — Tool System
Available tools for the ReAct agent:
  1. web_search   — DuckDuckGo free search
  2. calculator   — Safe math evaluation
  3. code_executor — Safe Python sandbox
"""

import ast
from typing import Dict, Any, Optional, Tuple, List
import math
import time
import signal
import subprocess
import tempfile
import os
import re
from typing import Dict, Any, Optional
from loguru import logger

try:
    from duckduckgo_search import DDGS
    DDG_AVAILABLE = True
except ImportError:
    DDG_AVAILABLE = False
    logger.warning("duckduckgo-search not installed — web search disabled")


# ─────────────────────────────────────────────
# Tool registry
# ─────────────────────────────────────────────
TOOL_DEFINITIONS = {
    "web_search": {
        "name": "web_search",
        "description": "Search the web for current information. Use for: news, facts, recent events.",
        "parameters": {"query": "string — the search query"},
        "example": "[TOOL: web_search] latest AI research 2024",
    },
    "calculator": {
        "name": "calculator",
        "description": "Evaluate mathematical expressions. Supports: +,-,*,/,**,sqrt,sin,cos,log,pi,e",
        "parameters": {"expression": "string — math expression to evaluate"},
        "example": "[TOOL: calculator] sqrt(144) + 2**10",
    },
    "code_executor": {
        "name": "code_executor",
        "description": "Execute Python code in a safe sandbox. Returns stdout output.",
        "parameters": {"code": "string — Python code to execute"},
        "example": "[TOOL: code_executor] print([x**2 for x in range(10)])",
    },
}


def get_tool_descriptions() -> str:
    """Format tool descriptions for LLM system prompt."""
    lines = ["Available Tools:"]
    for tool in TOOL_DEFINITIONS.values():
        lines.append(f"  • {tool['name']}: {tool['description']}")
        lines.append(f"    Usage: {tool['example']}")
    return "\n".join(lines)


# ─────────────────────────────────────────────
# Tool dispatcher
# ─────────────────────────────────────────────
def execute_tool(tool_name: str, argument: str) -> Dict[str, Any]:
    """
    Execute a named tool with its argument.

    Returns:
        {"success": bool, "result": str, "tool": str, "argument": str}
    """
    tool_name = tool_name.strip().lower()
    argument = argument.strip()

    logger.info(f"Tool call: [{tool_name}] {argument[:100]}")

    if tool_name == "web_search":
        return _web_search(argument)
    elif tool_name == "calculator":
        return _calculator(argument)
    elif tool_name == "code_executor":
        return _code_executor(argument)
    else:
        return {
            "success": False,
            "result": f"Unknown tool: {tool_name}. Available: {', '.join(TOOL_DEFINITIONS.keys())}",
            "tool": tool_name,
            "argument": argument,
        }


# ─────────────────────────────────────────────
# Web Search (DuckDuckGo — free, no API key)
# ─────────────────────────────────────────────
def _web_search(query: str, max_results: int = 5) -> Dict[str, Any]:
    """Search DuckDuckGo and return top results."""
    if not DDG_AVAILABLE:
        return {
            "success": False,
            "result": "Web search not available — install duckduckgo-search",
            "tool": "web_search",
            "argument": query,
        }

    try:
        results = []
        with DDGS() as ddgs:
            for r in ddgs.text(query, max_results=max_results):
                results.append({
                    "title": r.get("title", ""),
                    "snippet": r.get("body", ""),
                    "url": r.get("href", ""),
                })

        if not results:
            return {
                "success": True,
                "result": "No results found.",
                "tool": "web_search",
                "argument": query,
            }

        # Format as readable text
        formatted = f"Search results for: '{query}'\n\n"
        for i, r in enumerate(results, 1):
            formatted += f"{i}. **{r['title']}**\n"
            formatted += f"   {r['snippet']}\n"
            formatted += f"   Source: {r['url']}\n\n"

        return {
            "success": True,
            "result": formatted.strip(),
            "tool": "web_search",
            "argument": query,
            "raw_results": results,
        }

    except Exception as e:
        logger.error(f"Web search error: {e}")
        return {
            "success": False,
            "result": f"Search failed: {str(e)}",
            "tool": "web_search",
            "argument": query,
        }


# ─────────────────────────────────────────────
# Calculator (safe AST eval)
# ─────────────────────────────────────────────

# Allowed AST node types for safe math evaluation
SAFE_NODES = {
    ast.Expression, ast.BinOp, ast.UnaryOp, ast.Num, ast.Constant,
    ast.Add, ast.Sub, ast.Mult, ast.Div, ast.Pow, ast.FloorDiv, ast.Mod,
    ast.UAdd, ast.USub, ast.Call, ast.Name, ast.Load,
}

SAFE_FUNCTIONS = {
    "sqrt": math.sqrt, "sin": math.sin, "cos": math.cos, "tan": math.tan,
    "log": math.log, "log10": math.log10, "log2": math.log2,
    "abs": abs, "round": round, "ceil": math.ceil, "floor": math.floor,
    "exp": math.exp, "factorial": math.factorial,
    "pi": math.pi, "e": math.e, "inf": math.inf,
    "pow": pow, "min": min, "max": max, "sum": sum,
}


def _safe_eval(expr: str) -> Any:
    """Evaluate a math expression safely using AST."""
    tree = ast.parse(expr.strip(), mode="eval")

    # Validate all nodes
    for node in ast.walk(tree):
        if type(node) not in SAFE_NODES:
            raise ValueError(f"Unsafe operation: {type(node).__name__}")

    return eval(compile(tree, "<calc>", "eval"), {"__builtins__": {}}, SAFE_FUNCTIONS)


def _calculator(expression: str) -> Dict[str, Any]:
    """Evaluate a mathematical expression."""
    # Clean up the expression
    expression = expression.strip()
    # Replace common shortcuts
    expression = expression.replace("^", "**")  # ^ → power

    try:
        result = _safe_eval(expression)
        # Format result nicely
        if isinstance(result, float):
            if result == int(result) and abs(result) < 1e15:
                result_str = str(int(result))
            else:
                result_str = f"{result:.10g}"
        else:
            result_str = str(result)

        return {
            "success": True,
            "result": f"{expression} = {result_str}",
            "value": result,
            "tool": "calculator",
            "argument": expression,
        }

    except ZeroDivisionError:
        return {"success": False, "result": "Error: Division by zero", "tool": "calculator", "argument": expression}
    except (ValueError, TypeError) as e:
        return {"success": False, "result": f"Math error: {str(e)}", "tool": "calculator", "argument": expression}
    except Exception as e:
        return {"success": False, "result": f"Calculation failed: {str(e)}", "tool": "calculator", "argument": expression}


# ─────────────────────────────────────────────
# Code Executor (safe Python sandbox)
# ─────────────────────────────────────────────

BLOCKED_IMPORTS = {
    "os", "sys", "subprocess", "shutil", "socket", "urllib",
    "http", "ftplib", "smtplib", "importlib", "ctypes",
    "multiprocessing", "threading", "asyncio", "signal",
    "pickle", "shelve", "builtins",
}

BLOCKED_FUNCTIONS = {
    "exec", "eval", "compile", "__import__", "open",
    "globals", "locals", "vars", "dir", "getattr", "setattr",
    "delattr", "hasattr",
}


def _is_safe_code(code: str) -> Tuple[bool, str]:
    """Check code for dangerous patterns before execution."""
    # Block dangerous imports
    for blocked in BLOCKED_IMPORTS:
        pattern = rf"\b(import\s+{blocked}|from\s+{blocked}\b)"
        if re.search(pattern, code):
            return False, f"Import '{blocked}' is not allowed"

    # Block dangerous builtins
    for fn in BLOCKED_FUNCTIONS:
        if re.search(rf"\b{fn}\s*\(", code):
            return False, f"Function '{fn}' is not allowed"

    # Block file operations
    if re.search(r"\bopen\s*\(", code):
        return False, "File operations not allowed"

    return True, "OK"


def _code_executor(code: str, timeout: int = 10) -> Dict[str, Any]:
    """
    Execute Python code in a restricted subprocess sandbox.
    Has a strict timeout and output size limit.
    """
    safe, reason = _is_safe_code(code)
    if not safe:
        return {
            "success": False,
            "result": f"Code blocked for safety: {reason}",
            "tool": "code_executor",
            "argument": code,
        }

    # Write code to temp file
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", delete=False, dir="/tmp"
    ) as tmp:
        # Add safety wrapper
        safe_code = f"""
import math, json, re, datetime, random, itertools, functools, collections
from math import *

{code}
"""
        tmp.write(safe_code)
        tmp_path = tmp.name

    try:
        result = subprocess.run(
            ["python3", tmp_path],
            capture_output=True,
            text=True,
            timeout=timeout,
            # Restrict resources
            env={
                "PATH": "/usr/bin:/bin",
                "HOME": "/tmp",
            },
        )

        stdout = result.stdout[:2000]  # Limit output
        stderr = result.stderr[:500]

        if result.returncode == 0:
            output = stdout if stdout else "(no output)"
            return {
                "success": True,
                "result": f"Code output:\n{output}",
                "tool": "code_executor",
                "argument": code,
            }
        else:
            error = stderr or "Unknown error"
            # Clean up stack trace for readability
            error_lines = error.strip().split("\n")
            clean_error = error_lines[-1] if error_lines else error
            return {
                "success": False,
                "result": f"Code error: {clean_error}",
                "tool": "code_executor",
                "argument": code,
            }

    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "result": f"Code execution timed out ({timeout}s limit)",
            "tool": "code_executor",
            "argument": code,
        }
    except Exception as e:
        return {
            "success": False,
            "result": f"Execution error: {str(e)}",
            "tool": "code_executor",
            "argument": code,
        }
    finally:
        try:
            os.unlink(tmp_path)
        except Exception:
            pass