#!/usr/bin/env python3
"""
nima_agent_layer.py — OpenClaw-style self-authoring agent layer.

Nima can create, use, and compose tools. When she doesn't have a tool
she needs, she writes one, tests it in a sandbox, and saves it to her
toolbox for later.

NEUROBIOLOGICAL MAPPING:
  ToolRegistry  = basal ganglia + cerebellum (procedural memory)
  ToolExecutor  = primary motor cortex + spinal cord (action execution)
  ToolPlanner   = prefrontal cortex (executive planning)
  ToolCreator   = prefrontal ↔ motor integration (human-unique tool-making)
  ToolCombiner  = right hemisphere + PFC (creative synthesis)
  ToolSandbox   = cerebellum (safe practice space — test before using)

ENVIRONMENT VARIABLES:
  NIMA_TOOLS_DIR         # where tool files live (default: ./tools/)
  NIMA_SANDBOX_DIR       # where sandbox execution happens (default: ./sandbox/)
  NIMA_TOOL_AUTO_APPROVE # 1 = new tools don't need human approval
  NIMA_TOOL_TIMEOUT_S    # execution timeout (default: 30)
  NIMA_TOOL_MEMORY_MB    # memory limit (default: 256)
"""
from __future__ import annotations

import ast
import json
import logging
import os
import subprocess
import sys
import tempfile
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

logger = logging.getLogger("NimaAgent")

FORBIDDEN_MODULES = {
    "os.system", "subprocess.Popen", "subprocess.call", "subprocess.run",
    "os.popen", "commands.getoutput",
}
RESTRICTED_MODULES = {
    "eval", "exec", "compile",
    "os.remove", "os.rmdir", "os.unlink", "shutil.rmtree",
    "socket", "http.client", "urllib.request", "requests",
    "ctypes", "cffi",
}
ALLOWED_MODULES = {
    "math", "statistics", "random", "itertools", "collections",
    "json", "re", "string", "textwrap", "unicodedata",
    "datetime", "calendar", "hashlib", "base64",
    "fractions", "decimal", "numbers",
    "typing", "dataclasses", "enum",
    "functools", "operator", "pathlib",
    "csv", "io", "tempfile",
}


@dataclass
class ToolMetadata:
    tool_id: str
    name: str
    description: str
    function_signature: str
    file_path: str
    created_at: float = field(default_factory=time.time)
    created_by: str = "nima"
    use_count: int = 0
    last_used: float = 0.0
    approved: bool = False
    source_tools: List[str] = field(default_factory=list)
    tags: List[str] = field(default_factory=list)
    test_passed: bool = False

    def to_dict(self) -> Dict[str, Any]:
        return {k: v for k, v in self.__dict__.items()}


class ToolRegistry:
    """JSON-indexed library of tool files — Nima's procedural memory."""

    def __init__(self, tools_dir: Optional[str] = None) -> None:
        if tools_dir is None:
            tools_dir = os.environ.get("NIMA_TOOLS_DIR", os.path.join(os.getcwd(), "tools"))
        self.tools_dir = tools_dir
        self._tools: Dict[str, ToolMetadata] = {}
        self._lock = __import__("threading").Lock()
        os.makedirs(self.tools_dir, exist_ok=True)
        self._load()

    def _registry_path(self) -> str:
        return os.path.join(self.tools_dir, "_registry.json")

    def _load(self) -> None:
        path = self._registry_path()
        if not os.path.exists(path):
            return
        try:
            with open(path, "r") as f:
                data = json.load(f)
            for tid, md in data.items():
                self._tools[tid] = ToolMetadata(**md)
            logger.info("[ToolRegistry] loaded %d tools", len(self._tools))
        except Exception as e:
            logger.warning("[ToolRegistry] load failed: %s", e)

    def _save(self) -> None:
        path = self._registry_path()
        try:
            tmp = path + ".tmp"
            with open(tmp, "w") as f:
                json.dump({tid: t.to_dict() for tid, t in self._tools.items()},
                          f, indent=2, default=str)
            os.replace(tmp, path)
        except Exception as e:
            logger.warning("[ToolRegistry] save failed: %s", e)

    def register(self, metadata: ToolMetadata) -> None:
        with self._lock:
            self._tools[metadata.tool_id] = metadata
            self._save()

    def get(self, tool_id: str) -> Optional[ToolMetadata]:
        return self._tools.get(tool_id)

    def find_by_name(self, name: str) -> Optional[ToolMetadata]:
        for t in self._tools.values():
            if t.name == name:
                return t
        return None

    def find_by_capability(self, description_query: str) -> List[ToolMetadata]:
        query_lower = description_query.lower()
        query_words = set(query_lower.split())
        results = []
        for t in self._tools.values():
            desc_lower = t.description.lower()
            name_lower = t.name.lower()
            tags_lower = [tag.lower() for tag in t.tags]
            all_words = set()
            all_words.update(desc_lower.split())
            all_words.update(name_lower.replace("_", " ").split())
            for tag in tags_lower:
                all_words.update(tag.split())
            overlap = len(query_words & all_words)
            for qw in query_words:
                if len(qw) < 3:
                    continue
                if qw in name_lower or qw in desc_lower:
                    overlap += 1
                    continue
                for tag in tags_lower:
                    if qw in tag:
                        overlap += 1
                        break
                else:
                    for word in all_words:
                        if word.startswith(qw) and len(word) > len(qw):
                            overlap += 1
                            break
            if overlap > 0:
                results.append((t, overlap))
        results.sort(key=lambda x: -x[1])
        return [t for t, _ in results[:10]]

    def list_all(self) -> List[ToolMetadata]:
        return list(self._tools.values())

    def list_approved(self) -> List[ToolMetadata]:
        return [t for t in self._tools.values() if t.approved]

    def approve(self, tool_id: str) -> bool:
        with self._lock:
            t = self._tools.get(tool_id)
            if t is None:
                return False
            t.approved = True
            self._save()
        return True

    def record_use(self, tool_id: str) -> None:
        with self._lock:
            t = self._tools.get(tool_id)
            if t is None:
                return
            t.use_count += 1
            t.last_used = time.time()
            self._save()

    def get_stats(self) -> Dict[str, Any]:
        return {
            "total_tools": len(self._tools),
            "approved_tools": sum(1 for t in self._tools.values() if t.approved),
            "self_authored": sum(1 for t in self._tools.values() if t.created_by == "nima"),
            "total_uses": sum(t.use_count for t in self._tools.values()),
        }


class ToolSandbox:
    """Safe execution environment for tool code."""

    def __init__(self, sandbox_dir: Optional[str] = None) -> None:
        if sandbox_dir is None:
            sandbox_dir = os.environ.get("NIMA_SANDBOX_DIR", os.path.join(os.getcwd(), "sandbox"))
        self.sandbox_dir = sandbox_dir
        os.makedirs(self.sandbox_dir, exist_ok=True)
        self.timeout_s = float(os.environ.get("NIMA_TOOL_TIMEOUT_S", "30"))
        self.memory_limit_mb = float(os.environ.get("NIMA_TOOL_MEMORY_MB", "256"))

    def scan_code(self, code: str) -> Tuple[bool, List[str]]:
        issues: List[str] = []
        try:
            tree = ast.parse(code)
        except SyntaxError as e:
            return False, [f"SyntaxError: {e}"]
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    mod = alias.name
                    if mod in FORBIDDEN_MODULES:
                        issues.append(f"Forbidden import: {mod}")
                    elif mod in RESTRICTED_MODULES:
                        issues.append(f"Restricted import (needs approval): {mod}")
            elif isinstance(node, ast.ImportFrom):
                mod = node.module or ""
                if mod in FORBIDDEN_MODULES:
                    issues.append(f"Forbidden import: {mod}")
                elif mod in RESTRICTED_MODULES:
                    issues.append(f"Restricted import (needs approval): {mod}")
        blocked = [i for i in issues if "Forbidden" in i or "Restricted" in i]
        return len(blocked) == 0, issues

    def _build_runner_code(self, tool_file: str, function_name: str) -> str:
        """Build the runner script that executes a tool in a subprocess.

        Includes memory limiting via resource.setrlimit where supported.
        """
        memory_bytes = int(self.memory_limit_mb * 1024 * 1024)
        runner_code = (
            f"import sys, json, importlib.util, resource\n"
            f"# Enforce memory limit (Linux/macOS only)\n"
            f"try:\n"
            f"    resource.setrlimit(resource.RLIMIT_AS, ({memory_bytes}, {memory_bytes}))\n"
            f"except (ValueError, AttributeError):\n"
            f"    pass  # not available on this platform\n"
            f"spec = importlib.util.spec_from_file_location('tool', {tool_file!r})\n"
            f"mod = importlib.util.module_from_spec(spec)\n"
            f"spec.loader.exec_module(mod)\n"
            f"fn = getattr(mod, {function_name!r})\n"
            f"args = json.loads(sys.argv[1])\n"
            f"kwargs = json.loads(sys.argv[2])\n"
            f"try:\n"
            f"    result = fn(*args, **kwargs)\n"
            f"    print('__RESULT__' + json.dumps({{'result': result}}, default=str))\n"
            f"except Exception as e:\n"
            f"    print('__ERROR__' + json.dumps({{'error': str(e)}}))\n"
        )
        return runner_code

    def execute(self, code: str, function_name: str,
                args: List[Any] = None, kwargs: Dict[str, Any] = None,
                ) -> Dict[str, Any]:
        args = args or []
        kwargs = kwargs or {}
        t0 = time.time()
        passes, issues = self.scan_code(code)
        blocked = [i for i in issues if "Forbidden" in i or "Restricted" in i]
        if blocked:
            return {"success": False, "result": None, "stderr": "Blocked by safety scan",
                    "execution_time_ms": 0, "issues": blocked}
        tool_file = os.path.join(self.sandbox_dir, f"_tool_{uuid.uuid4().hex[:8]}.py")
        runner_file = os.path.join(self.sandbox_dir, f"_runner_{uuid.uuid4().hex[:8]}.py")
        try:
            with open(tool_file, "w") as f:
                f.write(code)
            runner_code = self._build_runner_code(tool_file, function_name)
            with open(runner_file, "w") as f:
                f.write(runner_code)
            proc = subprocess.run(
                [sys.executable, runner_file, json.dumps(args), json.dumps(kwargs)],
                capture_output=True, text=True, timeout=self.timeout_s, cwd=self.sandbox_dir,
            )
            elapsed_ms = (time.time() - t0) * 1000
            stdout, stderr = proc.stdout, proc.stderr
            result = None
            success = False
            for marker in ("__RESULT__", "__ERROR__"):
                if marker in stdout:
                    try:
                        payload = json.loads(stdout.split(marker)[-1].strip())
                        if marker == "__RESULT__":
                            result = payload.get("result")
                            success = True
                        else:
                            stderr = payload.get("error", stderr)
                    except Exception:
                        pass
                    break
            return {"success": success, "result": result,
                    "stdout": stdout[:2000], "stderr": stderr[:2000],
                    "execution_time_ms": round(elapsed_ms, 1), "issues": issues}
        except subprocess.TimeoutExpired:
            return {"success": False, "result": None, "stdout": "",
                    "stderr": f"Timeout after {self.timeout_s}s",
                    "execution_time_ms": round((time.time() - t0) * 1000, 1), "issues": issues}
        except Exception as e:
            return {"success": False, "result": None, "stdout": "",
                    "stderr": str(e), "execution_time_ms": 0, "issues": issues}
        finally:
            for p in (tool_file, runner_file):
                try:
                    os.remove(p)
                except Exception:
                    pass


class ToolExecutor:
    """Executes approved tools from the registry."""
    def __init__(self, registry: ToolRegistry, sandbox: ToolSandbox) -> None:
        self.registry = registry
        self.sandbox = sandbox

    def execute(self, tool_id: str, args: List[Any] = None,
                kwargs: Dict[str, Any] = None) -> Dict[str, Any]:
        tool = self.registry.get(tool_id)
        if tool is None:
            return {"success": False, "error": f"Tool not found: {tool_id}"}
        if not tool.approved:
            return {"success": False, "error": f"Tool not approved: {tool.name}"}
        tool_path = os.path.join(self.registry.tools_dir, tool.file_path)
        if not os.path.exists(tool_path):
            return {"success": False, "error": f"Tool file missing: {tool_path}"}
        with open(tool_path, "r") as f:
            code = f.read()
        result = self.sandbox.execute(code, tool.name, args, kwargs)
        if result["success"]:
            self.registry.record_use(tool_id)
        return result


class ToolPlanner:
    """Decides which tool to call for a given task."""
    def __init__(self, registry: ToolRegistry) -> None:
        self.registry = registry

    MAX_PLAN_TOOLS = 3  # maximum tools to return per plan

    def plan(self, task_description: str) -> Dict[str, Any]:
        candidates = self.registry.find_by_capability(task_description)
        approved = [t for t in candidates if t.approved]
        if not approved:
            return {"needs_tool": False, "tools": [],
                    "reasoning": "No approved tools match."}
        top_n = approved[:self.MAX_PLAN_TOOLS]
        return {"needs_tool": True, "tools": [{
            "tool_id": t.tool_id, "name": t.name,
            "args": [], "kwargs": {},
            "reason": f"Candidate for: {task_description[:80]}",
        } for t in top_n], "reasoning": f"Selected {len(top_n)} tool(s): {', '.join(t.name for t in top_n)}"}


class ToolCreator:
    """When Nima doesn't have a tool she needs, she writes one."""
    TOOL_TEMPLATE = '''"""Tool: {name} — {description}"""
import math, json, re
from typing import Any, Dict, List, Optional

def {function_name}({signature}):
    """{description}"""
    # TODO: Nima will fill this in via LLM
    pass
'''

    def __init__(self, registry: ToolRegistry, sandbox: ToolSandbox,
                 llm_provider: Optional[Any] = None) -> None:
        self.registry = registry
        self.sandbox = sandbox
        self.llm_provider = llm_provider
        self.auto_approve = bool(
            os.environ.get("NIMA_TOOL_AUTO_APPROVE", "0") in ("1", "true", "True"))

    def create_tool(self, name: str, description: str, function_signature: str,
                    test_cases: Optional[List[Dict[str, Any]]] = None) -> Dict[str, Any]:
        code = self._generate_code(name, description, function_signature, test_cases)
        test_results = []
        all_passed = True
        if test_cases:
            for i, tc in enumerate(test_cases):
                result = self.sandbox.execute(code, name, tc.get("args", []), tc.get("kwargs", {}))
                passed = result["success"]
                if passed and "expected" in tc:
                    passed = result["result"] == tc["expected"]
                test_results.append({"case": i, "passed": passed,
                                     "result": result.get("result"),
                                     "expected": tc.get("expected")})
                if not passed:
                    all_passed = False
        else:
            passes_scan, issues = self.sandbox.scan_code(code)
            all_passed = passes_scan
            test_results = [{"case": 0, "passed": passes_scan, "issues": issues}]
        approved = all_passed and self.auto_approve
        tool_id = f"tool_{uuid.uuid4().hex[:8]}"
        file_name = f"{name}.py"
        file_path = os.path.join(self.registry.tools_dir, file_name)
        with open(file_path, "w") as f:
            f.write(code)
        metadata = ToolMetadata(
            tool_id=tool_id, name=name, description=description,
            function_signature=f"{name}({function_signature})", file_path=file_name,
            created_by="nima", approved=approved, test_passed=all_passed,
            tags=self._extract_tags(description),
        )
        self.registry.register(metadata)
        return {"success": True, "tool_id": tool_id, "approved": approved,
                "test_results": test_results}

    def _generate_code(self, name, description, signature, test_cases=None) -> str:
        if self.llm_provider and getattr(self.llm_provider, "available", False):
            try:
                system_prompt = ("You are Nima's tool-creation subsystem. Write a single "
                                 "Python function. Use only safe stdlib modules. Output ONLY Python code.")
                user_prompt = f"Tool: {name}\nDescription: {description}\nSignature: def {name}({signature}):"
                code = self.llm_provider.generate(
                    system_prompt=system_prompt,
                    messages=[{"role": "user", "content": user_prompt}],
                    temperature=0.2, max_tokens=1024,
                )
                if code.startswith("```"):
                    lines = code.split("\n")[1:]
                    if lines and lines[-1].strip() == "```":
                        lines = lines[:-1]
                    code = "\n".join(lines)
                return code
            except Exception:
                pass
        return self.TOOL_TEMPLATE.format(
            name=name, description=description, function_name=name, signature=signature,
        )

    def _extract_tags(self, description: str) -> List[str]:
        tags = []
        desc_lower = description.lower()
        if any(w in desc_lower for w in ["search", "find", "lookup"]):
            tags.append("search")
        if any(w in desc_lower for w in ["file", "read", "write"]):
            tags.append("file")
        if any(w in desc_lower for w in ["math", "calculate", "compute", "convert"]):
            tags.append("math")
        if any(w in desc_lower for w in ["web", "url"]):
            tags.append("web")
        if any(w in desc_lower for w in ["time", "date"]):
            tags.append("time")
        return tags or ["general"]


class ToolCombiner:
    """Composes existing tools into composite tools."""
    def __init__(self, registry: ToolRegistry, creator: ToolCreator) -> None:
        self.registry = registry
        self.creator = creator

    def combine(self, name: str, description: str,
                source_tool_ids: List[str], combination_logic: str) -> Dict[str, Any]:
        source_tools = []
        for tid in source_tool_ids:
            t = self.registry.get(tid)
            if t is None:
                return {"success": False, "error": f"Source tool not found: {tid}"}
            source_tools.append(t)
        full_desc = f"{description}\nCombines: {', '.join(t.name for t in source_tools)}\nLogic: {combination_logic}"
        result = self.creator.create_tool(name=name, description=full_desc,
                                          function_signature="*args, **kwargs")
        if result.get("success"):
            tool = self.registry.get(result["tool_id"])
            if tool:
                tool.source_tools = source_tool_ids
                tool.tags.append("composite")
                self.registry.register(tool)
        return result


class AgentLayer:
    """Facade that ties together the full agent layer."""
    def __init__(self, tools_dir: Optional[str] = None,
                 sandbox_dir: Optional[str] = None,
                 llm_provider: Optional[Any] = None) -> None:
        self.registry = ToolRegistry(tools_dir)
        self.sandbox = ToolSandbox(sandbox_dir)
        self.executor = ToolExecutor(self.registry, self.sandbox)
        self.planner = ToolPlanner(self.registry)
        self.creator = ToolCreator(self.registry, self.sandbox, llm_provider)
        self.combiner = ToolCombiner(self.registry, self.creator)

    def register_starter_tools(self) -> None:
        starter_code = {
            "calculator": '''"""Tool: calculator — Evaluate a mathematical expression safely."""
import ast, operator as op
_OPS = {ast.Add: op.add, ast.Sub: op.sub, ast.Mult: op.mul, ast.Div: op.truediv,
        ast.Pow: op.pow, ast.Mod: op.mod, ast.USub: op.neg, ast.UAdd: op.pos, ast.FloorDiv: op.floordiv}
def calculator(expression: str):
    """Evaluate a mathematical expression. Returns float."""
    def _eval(node):
        if isinstance(node, ast.Num): return node.n
        if isinstance(node, ast.Constant): return node.value
        if isinstance(node, ast.BinOp): return _OPS[type(node.op)](_eval(node.left), _eval(node.right))
        if isinstance(node, ast.UnaryOp): return _OPS[type(node.op)](_eval(node.operand))
        raise ValueError(f"Unsupported: {type(node).__name__}")
    tree = ast.parse(expression, mode="eval")
    return float(_eval(tree.body))
''',
            "text_stats": '''"""Tool: text_stats — Compute statistics about text."""
import re
from collections import Counter
def text_stats(text: str):
    """Compute word count, char count, sentence count, avg word length, top words."""
    words = re.findall(r"\\b\\w+\\b", text.lower())
    sentences = [s for s in re.split(r"[.!?]+", text) if s.strip()]
    word_lens = [len(w) for w in words]
    return {"word_count": len(words), "char_count": len(text),
            "sentence_count": len(sentences),
            "avg_word_length": sum(word_lens)/max(1,len(word_lens)),
            "top_words": Counter(words).most_common(5)}
''',
            "time_now": '''"""Tool: time_now — Get the current time and date."""
import time
from datetime import datetime
def time_now():
    """Get current time and date. Returns dict with iso, date, time, weekday."""
    now = datetime.now()
    return {"iso": now.isoformat(), "unix": time.time(),
            "date": now.strftime("%Y-%m-%d"), "time": now.strftime("%H:%M:%S"),
            "weekday": now.strftime("%A")}
''',
        }
        for name, code in starter_code.items():
            if self.registry.find_by_name(name) is not None:
                continue
            file_path = os.path.join(self.registry.tools_dir, f"{name}.py")
            with open(file_path, "w") as f:
                f.write(code)
            result = self.sandbox.execute(code, name, [], {})
            desc = code.split("—")[1].split('"""')[0].strip() if "—" in code else name
            metadata = ToolMetadata(
                tool_id=f"tool_{uuid.uuid4().hex[:8]}", name=name,
                description=desc, function_signature=f"{name}()", file_path=f"{name}.py",
                created_by="human", approved=True, test_passed=result["success"],
                tags=["built-in", name],
            )
            self.registry.register(metadata)

    def get_stats(self) -> Dict[str, Any]:
        return {"registry": self.registry.get_stats(),
                "sandbox_timeout_s": self.sandbox.timeout_s,
                "auto_approve": self.creator.auto_approve}


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
    import tempfile
    tmp_tools = tempfile.mkdtemp(prefix="nima_tools_test_")
    tmp_sandbox = tempfile.mkdtemp(prefix="nima_sandbox_test_")
    os.environ["NIMA_TOOLS_DIR"] = tmp_tools
    os.environ["NIMA_SANDBOX_DIR"] = tmp_sandbox
    os.environ["NIMA_TOOL_AUTO_APPROVE"] = "1"
    agent = AgentLayer()
    agent.register_starter_tools()
    calc = agent.registry.find_by_name("calculator")
    if calc:
        r = agent.executor.execute(calc.tool_id, args=["2 + 3 * 4"])
        print(f"Calculator: 2 + 3 * 4 = {r.get('result')} (expected 14.0)")
    ts = agent.registry.find_by_name("text_stats")
    if ts:
        r = agent.executor.execute(ts.tool_id, args=["Hello world. Testing."])
        print(f"Text stats: {r.get('result')}")
    tn = agent.registry.find_by_name("time_now")
    if tn:
        r = agent.executor.execute(tn.tool_id, args=[])
        print(f"Time: {r.get('result', {}).get('iso', '?')}")
    print(f"Stats: {json.dumps(agent.get_stats(), indent=2)}")
    import shutil
    shutil.rmtree(tmp_tools, ignore_errors=True)
    shutil.rmtree(tmp_sandbox, ignore_errors=True)
    print("\nAgent layer test PASSED")