plane_mode_scholar/core/study_agent.py

"""Autonomous study agent — plans and executes sessions with minimal user input."""

from __future__ import annotations

import json
import logging
import time
from collections.abc import Iterator
from dataclasses import dataclass, field
from typing import Any

from plane_mode_scholar.core import llm
from plane_mode_scholar.core.config import DEFAULT_USER_ID
from plane_mode_scholar.core.orchestration import Orchestrator
from plane_mode_scholar.gradio_ui.citations import linkify_citations
from plane_mode_scholar.memory.scheduler import get_due_reviews, pick_study_focus
from plane_mode_scholar.study.planner import SessionPlanner
from plane_mode_scholar.study.quiz import QuizGenerator
from plane_mode_scholar.storage.sqlite_store import SQLiteStore

logger = logging.getLogger(__name__)

AGENT_TOOLS = [
    {"name": "ensure_pack", "description": "Load or create demo study materials"},
    {"name": "start_session", "description": "Open a timed study session with memory recall"},
    {"name": "explain_topic", "description": "Grounded explanation with [1] citations"},
    {"name": "run_quiz", "description": "Generate and present a quiz question"},
    {"name": "surface_memory", "description": "Show what the coach remembers about the learner"},
    {"name": "plan_next", "description": "Decide the next study action from mastery + SRS"},
]


@dataclass
class AgentEvent:
    phase: str
    message: str
    tool: str | None = None
    data: dict[str, Any] = field(default_factory=dict)
    cycle_id: int = 0

    def to_dict(self) -> dict[str, Any]:
        return {
            "phase": self.phase,
            "message": self.message,
            "tool": self.tool,
            "data": self.data,
            "cycle_id": self.cycle_id,
            "timestamp": time.time(),
        }


class StudyAgent:
    """Monitor → plan → act loop for autonomous study coaching."""

    def __init__(
        self,
        store: SQLiteStore | None = None,
        orchestrator: Orchestrator | None = None,
    ) -> None:
        self.store = store or SQLiteStore()
        self.orchestrator = orchestrator or Orchestrator(self.store)
        self.planner = SessionPlanner(self.store)
        self.quiz_gen = QuizGenerator(self.store)
        self._cycle = 0

    def _next_cycle(self) -> int:
        self._cycle += 1
        return self._cycle

    def _emit(self, phase: str, message: str, tool: str | None = None, **data: Any) -> dict:
        return AgentEvent(
            phase=phase,
            message=message,
            tool=tool,
            data=data,
            cycle_id=self._next_cycle(),
        ).to_dict()

    def plan_next_action(self, pack_id: str, user_id: str, session_id: str | None) -> dict:
        due = get_due_reviews(self.store, pack_id, user_id)
        weak = self.orchestrator.mastery.get_weak_topics(pack_id)
        focus = pick_study_focus(self.store, pack_id, user_id, weak)
        turns = self.store.list_turns(session_id) if session_id else []

        if not session_id:
            return {"action": "start_session", "topic": focus, "reason": "No active session"}
        if due:
            return {
                "action": "explain_topic",
                "topic": due[0].content[:80],
                "reason": f"{len(due)} review(s) due from last session",
            }
        if len(turns) < 2:
            return {
                "action": "explain_topic",
                "topic": focus,
                "reason": "Opening explanation for weak focus area",
            }
        if len(turns) < 6:
            return {
                "action": "run_quiz",
                "topic": focus,
                "reason": "Check understanding before moving on",
            }
        return {
            "action": "surface_memory",
            "topic": focus,
            "reason": "Session depth reached — show memory + suggest wrap-up",
        }

    def ensure_pack(self, user_id: str, pack_id: str | None = None) -> str:
        if pack_id:
            pack = self.store.get_pack(pack_id)
            if pack:
                return pack_id
        packs = self.store.list_packs(user_id=user_id)
        if packs:
            return packs[0].pack_id
        from plane_mode_scholar.gradio_ui.layout import load_demo_pack

        load_demo_pack(user_id)
        packs = self.store.list_packs(user_id=user_id)
        if not packs:
            raise RuntimeError("Could not create demo study pack")
        return packs[0].pack_id

    def start_session(
        self, pack_id: str, user_id: str, goals: list[str] | None = None
    ) -> dict[str, Any]:
        result = self.planner.start_session(
            pack_id, goals=goals or ["Review key exam topics"], duration_minutes=45, user_id=user_id
        )
        return result

    def run_autopilot(
        self,
        user_id: str = DEFAULT_USER_ID,
        pack_id: str | None = None,
        max_steps: int = 5,
    ) -> Iterator[dict[str, Any]]:
        """One-click autonomous study flow — yields telemetry for SwarmGrid-style UI."""
        yield self._emit("BOOT", "Plane Mode Scholar agent online", tool="ensure_pack")

        try:
            pack_id = self.ensure_pack(user_id, pack_id)
            pack = self.store.get_pack(pack_id)
            pack_name = pack.name if pack else pack_id
            yield self._emit(
                "PACK",
                f"Study pack ready: {pack_name}",
                tool="ensure_pack",
                pack_id=pack_id,
            )

            session_result = self.start_session(pack_id, user_id)
            session_id = session_result["session"]["session_id"]
            due_count = session_result.get("due_count", 0)
            recall = session_result.get("retrieved_memories", [])
            yield self._emit(
                "SESSION",
                f"Session started — {due_count} due review(s), {len(recall)} memories loaded",
                tool="start_session",
                session_id=session_id,
                due_count=due_count,
                recall_preview=[m.get("content", "")[:80] for m in recall[:3]],
                recommended=session_result.get("recommended_start", ""),
            )

            steps_done = 0
            while steps_done < max_steps:
                plan = self.plan_next_action(pack_id, user_id, session_id)
                action = plan["action"]
                topic = plan.get("topic", "")
                yield self._emit(
                    "PLAN",
                    f"Next: {action} — {plan.get('reason', '')}",
                    tool="plan_next",
                    plan=plan,
                )

                if action == "explain_topic":
                    query = f"Explain {topic} clearly using my course materials"
                    prepared = self.orchestrator.prepare_turn(
                        query, pack_id, session_id, topic, user_id=user_id
                    )
                    yield self._emit(
                        "RETRIEVE",
                        f"Retrieved {len(prepared.chunks)} chunks, {len(prepared.memories)} memories",
                        tool="explain_topic",
                        chunks=len(prepared.chunks),
                        memories=len(prepared.memories),
                    )
                    accumulated = ""
                    t0 = time.time()
                    for partial in llm.generate_stream(prepared.prompt):
                        accumulated = partial
                        yield self._emit(
                            "STREAM",
                            accumulated,
                            tool="explain_topic",
                            delta=partial,
                            topic=topic,
                        )
                    result = self.orchestrator.finalize_turn(prepared, accumulated)
                    cited = linkify_citations(accumulated, prepared.chunks)
                    yield self._emit(
                        "EXPLAIN",
                        cited,
                        tool="explain_topic",
                        response=cited,
                        memory_writes=len(result.memory_writes),
                        inference_ms=int((time.time() - t0) * 1000),
                    )

                elif action == "run_quiz":
                    questions = self.quiz_gen.generate_quiz(
                        pack_id,
                        num_questions=1,
                        topic=topic,
                        llm_generate=llm.generate,
                        user_id=user_id,
                    )
                    if questions:
                        q = questions[0]
                        yield self._emit(
                            "QUIZ",
                            q.get("question", "Quiz ready"),
                            tool="run_quiz",
                            question=q,
                            topic=topic,
                        )
                    else:
                        yield self._emit("QUIZ", "Could not generate quiz", tool="run_quiz")

                elif action == "surface_memory":
                    active = self.store.list_memories(
                        pack_id=pack_id, user_id=user_id, status="active"
                    )
                    yield self._emit(
                        "MEMORY",
                        f"Coach remembers {len(active)} active memories",
                        tool="surface_memory",
                        memories=[m.to_dict() for m in active[:6]],
                    )

                steps_done += 1

            yield self._emit(
                "DONE",
                "Autopilot segment complete — continue chatting or tap Fly again",
                tool=None,
                session_id=session_id,
                pack_id=pack_id,
            )
        except Exception as e:
            logger.exception("Autopilot failed")
            yield self._emit("ERROR", str(e), tool=None)

    def agent_plan_json(self, pack_id: str, user_id: str, context: str) -> dict:
        """Optional LLM planner for Nemotron tool-calling style decisions."""
        prompt = f"""You are a study coach agent. Pick the single best next action.

Tools: {json.dumps([t["name"] for t in AGENT_TOOLS])}

Context:
{context}

Pack has materials indexed. Reply JSON only: {{"action": "<tool>", "topic": "<topic>", "reason": "<short>"}}"""
        raw = llm.generate(prompt, max_new_tokens=120)
        try:
            start = raw.find("{")
            end = raw.rfind("}") + 1
            if start >= 0 and end > start:
                return json.loads(raw[start:end])
        except json.JSONDecodeError:
            pass
        return self.plan_next_action(pack_id, user_id, None)