#!/usr/bin/env python3
"""
Vitalis Base Agent
Self-directed execution loop with memory, tools, and self-correction.
"""
import time
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional
from pathlib import Path
import json

class BaseTool:
    name: str        = "base_tool"
    description: str = "Override this"

    def run(self, **kwargs) -> Any:
        raise NotImplementedError

class AgentMemory:
    def __init__(self, path: Optional[Path] = None):
        self.path = path or (Path.home() / ".vitalis_workspace" / "agent_memory.json")
        self._store: List[Dict] = self._load()

    def _load(self) -> List[Dict]:
        if self.path.exists():
            try:
                return json.loads(self.path.read_text())
            except Exception:
                return []
        return []

    def add(self, role: str, content: str):
        self._store.append({"role": role, "content": content, "ts": time.time()})
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.path.write_text(json.dumps(self._store[-200:], indent=2))

    def recent(self, n: int = 10) -> List[Dict]:
        return self._store[-n:]

    def clear(self):
        self._store = []
        if self.path.exists():
            self.path.unlink()

class BaseAgent(ABC):
    MAX_STEPS     = 20
    STEP_DELAY    = 0.5

    def __init__(self, name: str = "VitalisAgent"):
        self.name    = name
        self.memory  = AgentMemory()
        self.tools: Dict[str, BaseTool] = {}
        self._steps  = 0
        self._running = False

    def register_tool(self, tool: BaseTool):
        self.tools[tool.name] = tool
        print(f"[AGENT:{self.name}] Tool registered: {tool.name}")

    def use_tool(self, tool_name: str, **kwargs) -> Any:
        tool = self.tools.get(tool_name)
        if not tool:
            return f"ERROR: Tool '{tool_name}' not found. Available: {list(self.tools.keys())}"
        try:
            result = tool.run(**kwargs)
            self.memory.add("tool", f"{tool_name}({kwargs}) → {str(result)[:200]}")
            return result
        except Exception as e:
            err = f"Tool '{tool_name}' failed: {e}"
            self.memory.add("error", err)
            return err

    @abstractmethod
    def think(self, state: Dict) -> Dict:
        """Override: decide next action given current state."""
        pass

    @abstractmethod
    def act(self, decision: Dict) -> Any:
        """Override: execute the decided action."""
        pass

    @abstractmethod
    def is_done(self, state: Dict) -> bool:
        """Override: return True when goal is reached."""
        pass

    def run(self, goal: str, initial_state: Optional[Dict] = None) -> Any:
        print(f"[AGENT:{self.name}] Starting. Goal: {goal}")
        self.memory.add("system", f"Goal: {goal}")
        state = initial_state or {"goal": goal, "step": 0, "results": []}
        self._running  = True
        self._steps    = 0
        last_result    = None

        while self._running and self._steps < self.MAX_STEPS:
            if self.is_done(state):
                print(f"[AGENT:{self.name}] Goal reached in {self._steps} steps.")
                break

            try:
                decision   = self.think(state)
                result     = self.act(decision)
                last_result = result
                state["step"]    = self._steps
                state["results"].append({"step": self._steps, "decision": decision, "result": str(result)[:300]})
                self.memory.add("assistant", f"Step {self._steps}: {decision} → {str(result)[:150]}")
            except Exception as e:
                print(f"[AGENT:{self.name}] Step {self._steps} error: {e}")
                self.memory.add("error", str(e))
                state["last_error"] = str(e)

            self._steps += 1
            time.sleep(self.STEP_DELAY)

        if self._steps >= self.MAX_STEPS:
            print(f"[AGENT:{self.name}] Max steps reached ({self.MAX_STEPS}).")

        self._running = False
        return last_result

    def stop(self):
        self._running = False
        print(f"[AGENT:{self.name}] Stopped.")