splitbit_llm/memory/goal_memory.py

"""Goal & Planning Memory — long-term persistent goals with execution plans.

Goals can be:
- Created with description, priority, deadline
- Broken into steps by the planner agent
- Tracked: pending → planning → in_progress → reviewing → completed/failed
- Executed step-by-step by agents
- Paused/resumed
- Linked via dependencies (one goal blocks another)
- Persisted to SQLite — survives restarts

Planning memory: stores generated plans, replanning history, and
step-by-step progress so work can resume after restarts.
"""

from __future__ import annotations

import hashlib
import json
import logging
import os
import sqlite3
import time
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger(__name__)

GOAL_STATUSES = ("pending", "planning", "in_progress", "reviewing", "completed", "failed", "paused")
GOAL_PRIORITIES = ("critical", "high", "medium", "low")


@dataclass
class Goal:
    """A long-term goal with execution plan."""
    id: str
    title: str
    description: str
    priority: str = "medium"
    status: str = "pending"
    parent_goal_id: str | None = None
    sub_goal_ids: list[str] = field(default_factory=list)
    steps: list[dict[str, Any]] = field(default_factory=list)
    current_step: int = 0
    created_at: float = 0.0
    updated_at: float = 0.0
    deadline: float | None = None
    completed_at: float | None = None
    tags: list[str] = field(default_factory=list)
    assigned_agent: str = ""  # which agent is working on this
    blocks_goal_ids: list[str] = field(default_factory=list)
    blocked_by_goal_ids: list[str] = field(default_factory=list)
    replan_count: int = 0
    metadata: dict[str, Any] = field(default_factory=dict)

    def as_dict(self) -> dict[str, Any]:
        return {
            "id": self.id, "title": self.title, "description": self.description,
            "priority": self.priority, "status": self.status,
            "parent_goal_id": self.parent_goal_id, "sub_goal_ids": self.sub_goal_ids,
            "steps": self.steps, "current_step": self.current_step,
            "created_at": self.created_at, "updated_at": self.updated_at,
            "deadline": self.deadline, "completed_at": self.completed_at,
            "tags": self.tags, "assigned_agent": self.assigned_agent,
            "blocks_goal_ids": self.blocks_goal_ids,
            "blocked_by_goal_ids": self.blocked_by_goal_ids,
            "replan_count": self.replan_count,
            "progress": self.progress(),
            "metadata": self.metadata,
        }

    def progress(self) -> float:
        if not self.steps:
            return 0.0 if self.status != "completed" else 1.0
        completed = sum(1 for s in self.steps if s.get("status") == "completed")
        return completed / len(self.steps)


@dataclass
class PlanHistory:
    """History of planning/replanning for a goal."""
    id: str
    goal_id: str
    plan_type: str  # "initial", "replan"
    steps: list[dict[str, Any]] = field(default_factory=list)
    reason: str = ""
    timestamp: float = field(default_factory=time.time)


class GoalMemory:
    """Long-term persistent goal and planning memory.

    SQLite-backed. Goals survive restarts. Plans are stored and
    can be resumed. Includes:
    - Goal CRUD
    - Planning (generate steps from description)
    - Step tracking (mark steps as completed/failed)
    - Dependencies (goal blocking)
    - Plan history (track replans)
    - Agent assignment (which agent is working on what)
    """

    PLANNING_PROMPT = """You are a goal planning system. Given a goal, create a detailed step-by-step execution plan.

Return JSON:
{
  "steps": [
    {"title": "<step title>", "description": "<what to do>", "tool": "<tool to use>", "estimated_time_s": 30},
    ...
  ],
  "sub_goals": [
    {"title": "<sub-goal title>", "description": "<sub-goal description>", "priority": "medium"}
  ]
}

Break the goal into 3-8 concrete, actionable steps. Each step should be independently executable."""

    def __init__(self, db_path: str) -> None:
        self.db_path = db_path
        os.makedirs(os.path.dirname(db_path) or ".", exist_ok=True)
        self._init_db()

    def _init_db(self) -> None:
        with sqlite3.connect(self.db_path) as conn:
            conn.executescript("""
                CREATE TABLE IF NOT EXISTS goals (
                    id TEXT PRIMARY KEY,
                    title TEXT NOT NULL,
                    description TEXT NOT NULL,
                    priority TEXT DEFAULT 'medium',
                    status TEXT DEFAULT 'pending',
                    parent_goal_id TEXT,
                    sub_goal_ids TEXT DEFAULT '[]',
                    steps TEXT DEFAULT '[]',
                    current_step INTEGER DEFAULT 0,
                    created_at REAL NOT NULL,
                    updated_at REAL NOT NULL,
                    deadline REAL,
                    completed_at REAL,
                    tags TEXT DEFAULT '[]',
                    assigned_agent TEXT DEFAULT '',
                    blocks_goal_ids TEXT DEFAULT '[]',
                    blocked_by_goal_ids TEXT DEFAULT '[]',
                    replan_count INTEGER DEFAULT 0,
                    metadata TEXT DEFAULT '{}'
                );
                CREATE INDEX IF NOT EXISTS idx_goals_status ON goals(status);
                CREATE INDEX IF NOT EXISTS idx_goals_priority ON goals(priority);
                CREATE INDEX IF NOT EXISTS idx_goals_agent ON goals(assigned_agent);

                CREATE TABLE IF NOT EXISTS plan_history (
                    id TEXT PRIMARY KEY,
                    goal_id TEXT,
                    plan_type TEXT,
                    steps TEXT DEFAULT '[]',
                    reason TEXT,
                    timestamp REAL
                );
                CREATE INDEX IF NOT EXISTS idx_plan_goal ON plan_history(goal_id);
            """)

    def create_goal(self, title: str, description: str, priority: str = "medium",
                    deadline: float | None = None, parent_goal_id: str | None = None,
                    tags: list[str] | None = None, assigned_agent: str = "",
                    blocks: list[str] | None = None, blocked_by: list[str] | None = None) -> Goal:
        """Create a new goal."""
        goal_id = hashlib.sha256(f"{title}:{time.time()}".encode()).hexdigest()[:16]
        now = time.time()
        goal = Goal(
            id=goal_id, title=title, description=description, priority=priority,
            status="pending", parent_goal_id=parent_goal_id,
            created_at=now, updated_at=now, deadline=deadline,
            tags=tags or [], assigned_agent=assigned_agent,
            blocks_goal_ids=blocks or [], blocked_by_goal_ids=blocked_by or [],
        )
        self._store(goal)

        if parent_goal_id:
            parent = self.get_goal(parent_goal_id)
            if parent:
                parent.sub_goal_ids.append(goal_id)
                self._store(parent)

        logger.info("Created goal '%s' (%s, priority: %s, agent: %s)", title, goal_id, priority, assigned_agent)
        return goal

    def plan_goal(self, goal_id: str, steps: list[dict[str, Any]] | None = None,
                  sub_goals: list[dict[str, Any]] | None = None,
                  plan_type: str = "initial", reason: str = "") -> dict[str, Any]:
        """Plan a goal — either with provided steps or mark for LLM planning.

        If steps are provided, they're stored directly.
        If not, the goal is marked as "planning" and the planner agent will generate steps.
        """
        goal = self.get_goal(goal_id)
        if goal is None:
            return {"success": False, "message": "Goal not found"}

        if steps:
            goal.steps = []
            for i, step in enumerate(steps):
                goal.steps.append({
                    "index": i,
                    "title": step.get("title", f"Step {i+1}"),
                    "description": step.get("description", ""),
                    "tool": step.get("tool", ""),
                    "estimated_time_s": step.get("estimated_time_s", 30),
                    "status": "pending",
                    "result": "",
                    "assigned_agent": step.get("agent", ""),
                })
            goal.current_step = 0
            goal.status = "in_progress"
            goal.updated_at = time.time()

            if plan_type == "replan":
                goal.replan_count += 1

            self._store(goal)

            # Store plan history
            plan_id = hashlib.sha256(f"{goal_id}:{time.time()}".encode()).hexdigest()[:16]
            with sqlite3.connect(self.db_path) as conn:
                conn.execute(
                    "INSERT INTO plan_history VALUES (?,?,?,?,?,?)",
                    (plan_id, goal_id, plan_type, json.dumps(steps), reason, time.time())
                )

            # Create sub-goals
            if sub_goals:
                for sg in sub_goals:
                    sub = self.create_goal(
                        title=sg.get("title", ""), description=sg.get("description", ""),
                        priority=sg.get("priority", "medium"), parent_goal_id=goal_id,
                    )
                    goal.sub_goal_ids.append(sub.id)
                self._store(goal)

            logger.info("Planned goal '%s': %d steps, %d sub-goals", goal.title, len(goal.steps), len(sub_goals or []))
            return {"success": True, "goal": goal.as_dict()}
        else:
            # Mark for LLM planning
            goal.status = "planning"
            self._store(goal)
            return {"success": True, "message": "Goal marked for planning", "goal": goal.as_dict()}

    def execute_step(self, goal_id: str, result: str, success: bool = True) -> dict[str, Any]:
        """Mark the current step as completed/failed and advance."""
        goal = self.get_goal(goal_id)
        if goal is None:
            return {"success": False, "message": "Goal not found"}
        if goal.status not in ("in_progress", "reviewing"):
            return {"success": False, "message": f"Goal is {goal.status}"}
        if goal.current_step >= len(goal.steps):
            goal.status = "completed"
            goal.completed_at = time.time()
            self._store(goal)
            return {"success": True, "message": "Goal already completed", "goal": goal.as_dict()}

        step = goal.steps[goal.current_step]
        step["status"] = "completed" if success else "failed"
        step["result"] = result
        step["completed_at"] = time.time()

        goal.current_step += 1
        goal.updated_at = time.time()

        if goal.current_step >= len(goal.steps):
            goal.status = "completed"
            goal.completed_at = time.time()

        self._store(goal)

        logger.info("Step %d/%d for '%s': %s", goal.current_step, len(goal.steps), goal.title, step["status"])
        return {"success": True, "step": step, "goal": goal.as_dict()}

    def get_goal(self, goal_id: str) -> Goal | None:
        with sqlite3.connect(self.db_path) as conn:
            row = conn.execute("SELECT * FROM goals WHERE id = ?", (goal_id,)).fetchone()
        return self._row_to_goal(row) if row else None

    def list_goals(self, status: str | None = None, priority: str | None = None,
                   agent: str | None = None) -> list[Goal]:
        query = "SELECT * FROM goals"
        params: list[Any] = []
        conditions: list[str] = []
        if status:
            conditions.append("status = ?")
            params.append(status)
        if priority:
            conditions.append("priority = ?")
            params.append(priority)
        if agent:
            conditions.append("assigned_agent = ?")
            params.append(agent)
        if conditions:
            query += " WHERE " + " AND ".join(conditions)
        query += " ORDER BY CASE priority WHEN 'critical' THEN 0 WHEN 'high' THEN 1 WHEN 'medium' THEN 2 ELSE 3 END, created_at DESC"
        with sqlite3.connect(self.db_path) as conn:
            rows = conn.execute(query, params).fetchall()
        return [self._row_to_goal(r) for r in rows]

    def update_goal(self, goal_id: str, **kwargs: Any) -> Goal | None:
        goal = self.get_goal(goal_id)
        if goal is None:
            return None
        for k, v in kwargs.items():
            if hasattr(goal, k):
                setattr(goal, k, v)
        goal.updated_at = time.time()
        self._store(goal)
        return goal

    def delete_goal(self, goal_id: str) -> bool:
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute("DELETE FROM goals WHERE id = ?", (goal_id,))
            return cursor.rowcount > 0

    def get_active_goals(self) -> list[Goal]:
        return self.list_goals(status="in_progress") + self.list_goals(status="planning")

    def get_pending_goals(self) -> list[Goal]:
        return self.list_goals(status="pending")

    def get_goals_for_agent(self, agent_name: str) -> list[Goal]:
        return self.list_goals(agent=agent_name)

    def assign_agent(self, goal_id: str, agent_name: str) -> Goal | None:
        return self.update_goal(goal_id, assigned_agent=agent_name, status="in_progress")

    def get_goal_context(self) -> str:
        """Get active goals as context for the LLM."""
        active = self.get_active_goals()
        if not active:
            return ""
        lines = ["Active Goals:"]
        for g in active:
            progress = g.progress()
            lines.append(f"- [{g.priority}] {g.title} ({g.status}, {progress:.0%} done, agent: {g.assigned_agent})")
            if g.current_step < len(g.steps):
                step = g.steps[g.current_step]
                lines.append(f"  Next: {step['title']}")
        return "\n".join(lines)

    def get_plan_history(self, goal_id: str) -> list[PlanHistory]:
        with sqlite3.connect(self.db_path) as conn:
            rows = conn.execute(
                "SELECT * FROM plan_history WHERE goal_id = ? ORDER BY timestamp DESC",
                (goal_id,)
            ).fetchall()
        return [PlanHistory(
            id=r[0], goal_id=r[1], plan_type=r[2],
            steps=json.loads(r[3]), reason=r[4], timestamp=r[5]
        ) for r in rows]

    def _store(self, goal: Goal) -> None:
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                """INSERT OR REPLACE INTO goals
                   (id, title, description, priority, status, parent_goal_id, sub_goal_ids,
                    steps, current_step, created_at, updated_at, deadline, completed_at,
                    tags, assigned_agent, blocks_goal_ids, blocked_by_goal_ids, replan_count, metadata)
                   VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)""",
                (goal.id, goal.title, goal.description, goal.priority, goal.status,
                 goal.parent_goal_id, json.dumps(goal.sub_goal_ids), json.dumps(goal.steps),
                 goal.current_step, goal.created_at, goal.updated_at, goal.deadline,
                 goal.completed_at, json.dumps(goal.tags), goal.assigned_agent,
                 json.dumps(goal.blocks_goal_ids), json.dumps(goal.blocked_by_goal_ids),
                 goal.replan_count, json.dumps(goal.metadata)),
            )

    def _row_to_goal(self, row: tuple) -> Goal:
        return Goal(
            id=row[0], title=row[1], description=row[2], priority=row[3], status=row[4],
            parent_goal_id=row[5], sub_goal_ids=json.loads(row[6]), steps=json.loads(row[7]),
            current_step=row[8], created_at=row[9], updated_at=row[10], deadline=row[11],
            completed_at=row[12], tags=json.loads(row[13]), assigned_agent=row[14],
            blocks_goal_ids=json.loads(row[15]), blocked_by_goal_ids=json.loads(row[16]),
            replan_count=row[17], metadata=json.loads(row[18]),
        )

    def get_stats(self) -> dict[str, Any]:
        with sqlite3.connect(self.db_path) as conn:
            total = conn.execute("SELECT COUNT(*) FROM goals").fetchone()[0]
            active = conn.execute("SELECT COUNT(*) FROM goals WHERE status IN ('in_progress','planning')").fetchone()[0]
            completed = conn.execute("SELECT COUNT(*) FROM goals WHERE status = 'completed'").fetchone()[0]
            failed = conn.execute("SELECT COUNT(*) FROM goals WHERE status = 'failed'").fetchone()[0]
            pending = conn.execute("SELECT COUNT(*) FROM goals WHERE status = 'pending'").fetchone()[0]
            plans = conn.execute("SELECT COUNT(*) FROM plan_history").fetchone()[0]
        return {
            "total": total, "active": active, "completed": completed,
            "failed": failed, "pending": pending, "plan_history": plans,
        }