app/services/agent_control.py

from __future__ import annotations

import json
import os
import time
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal

from app.schemas.commit_adaptation import CommitAdaptationReview

if TYPE_CHECKING:
    from app.agents.cerebras_client import CerebrasClient

_DEFAULT_ROOT = os.path.expanduser("~/.studybuddy/agent_control")
_COMMIT_REVIEWER_SKILL_PATH = (
    Path(__file__).resolve().parents[1] / "memory_skills" / "research-memory-reviewer" / "SKILL.md"
)

ProfileScope = Literal["student", "project"]
ProfileSource = Literal["manual", "inferred"]

BASE_CONTRACT = (
    "ResearchMate is local-first and grounded in the student's uploaded project material. "
    "It may adapt style and workflow, but it must not invent source facts, overwrite manual "
    "preferences, or hide memory failures from the project status surfaces."
)

DEFAULT_AGENT_SKILLS = {
    "brain": "Organize project material into a research map shaped by project memory.",
    "tutor": "Explain from grounded evidence, adapting order and rigor to the student profile.",
    "pair_buddy": "Help the student think aloud, ask grounded follow-ups, and preserve developing ideas.",
    "recommender": "Recommend papers from uploaded project papers, verified citation context, and student interests.",
    "visualization": (
        "Research missing evidence, expose assumptions and provenance, and never present illustrative values "
        "as observed data or trained-model weights."
    ),
}


def _profile_record() -> dict[str, Any]:
    return {"manual_text": "", "inferred_text": "", "updated_at": 0.0}


class AgentControlService:
    def __init__(self, root: str | None = None, client: "CerebrasClient | None" = None) -> None:
        self.root = root or _DEFAULT_ROOT
        os.makedirs(self.root, exist_ok=True)
        self._client = client

    def _path(self, project_id: str) -> str:
        safe = (project_id or "global").replace("/", "_").replace("\\", "_").replace("..", "_")
        return os.path.join(self.root, f"{safe}.json")

    def _default_state(self, project_id: str) -> dict[str, Any]:
        return {
            "project_id": project_id,
            "student": _profile_record(),
            "project": _profile_record(),
            "skills": {agent_id: _profile_record() | {"default_text": text} for agent_id, text in DEFAULT_AGENT_SKILLS.items()},
            "evolution": [],
            "updated_at": time.time(),
        }

    def _load(self, project_id: str) -> dict[str, Any]:
        path = self._path(project_id)
        if not os.path.exists(path):
            return self._default_state(project_id)
        try:
            with open(path, encoding="utf-8") as f:
                state = self._default_state(project_id) | json.load(f)
            state.setdefault("skills", {})
            for agent_id, text in DEFAULT_AGENT_SKILLS.items():
                state["skills"].setdefault(agent_id, _profile_record() | {"default_text": text})
            return state
        except Exception:
            return self._default_state(project_id)

    def _save(self, project_id: str, state: dict[str, Any]) -> dict[str, Any]:
        state["updated_at"] = time.time()
        os.makedirs(self.root, exist_ok=True)
        with open(self._path(project_id), "w", encoding="utf-8") as f:
            json.dump(state, f, indent=2)
        return state

    def get_profile(self, project_id: str) -> dict[str, Any]:
        state = self._load(project_id)
        return {
            "project_id": project_id,
            "student": state["student"],
            "project": state["project"],
            "updated_at": state["updated_at"],
        }

    def update_profile(self, project_id: str, scope: ProfileScope, text: str, source: ProfileSource = "manual") -> dict[str, Any]:
        state = self._load(project_id)
        key = "manual_text" if source == "manual" else "inferred_text"
        state[scope][key] = text
        state[scope]["updated_at"] = time.time()
        self._save(project_id, state)
        return self.get_profile(project_id)

    def get_skills(self, project_id: str) -> dict[str, Any]:
        state = self._load(project_id)
        return {"project_id": project_id, "skills": state["skills"]}

    def update_skill(self, project_id: str, skill_id: str, text: str, source: ProfileSource = "manual") -> dict[str, Any]:
        state = self._load(project_id)
        state["skills"].setdefault(skill_id, _profile_record() | {"default_text": ""})
        key = "manual_text" if source == "manual" else "inferred_text"
        state["skills"][skill_id][key] = text
        state["skills"][skill_id]["updated_at"] = time.time()
        self._save(project_id, state)
        return self.get_skills(project_id)

    def record_evolution(self, project_id: str, event_type: str, message: str, metadata: dict[str, Any] | None = None) -> dict[str, Any]:
        state = self._load(project_id)
        event = {
            "event_type": event_type,
            "message": message,
            "metadata": metadata or {},
            "created_at": time.time(),
        }
        state["evolution"].append(event)
        self._save(project_id, state)
        return event

    def get_evolution(self, project_id: str) -> list[dict[str, Any]]:
        state = self._load(project_id)
        return sorted(state["evolution"], key=lambda row: row.get("created_at", 0), reverse=True)

    @staticmethod
    def _effective(record: dict[str, Any]) -> str:
        return record.get("manual_text") or record.get("inferred_text") or record.get("default_text", "")

    def preview(self, project_id: str, agent_id: str = "tutor") -> dict[str, Any]:
        state = self._load(project_id)
        skill = state["skills"].get(agent_id, _profile_record() | {"default_text": DEFAULT_AGENT_SKILLS.get(agent_id, "")})
        layers = [
            {"name": "Base Contract", "content": BASE_CONTRACT, "editable": False},
            {"name": "Student Profile", "content": self._effective(state["student"]), "editable": True},
            {"name": "Project Profile", "content": self._effective(state["project"]), "editable": True},
            {"name": f"{agent_id} Skill", "content": self._effective(skill), "editable": True},
        ]
        return {
            "project_id": project_id,
            "agent_id": agent_id,
            "layers": layers,
            "composed_prompt": "\n\n".join(layer["content"] for layer in layers if layer["content"]),
        }

    def reset(self, project_id: str, scope: str = "all") -> dict[str, Any]:
        state = self._load(project_id)
        if scope in {"all", "profiles"}:
            state["student"] = _profile_record()
            state["project"] = _profile_record()
        if scope in {"all", "skills"}:
            for agent_id, text in DEFAULT_AGENT_SKILLS.items():
                state["skills"][agent_id] = _profile_record() | {"default_text": text}
        if scope == "all":
            state["evolution"] = []
        self._save(project_id, state)
        return {"ok": True, "scope": scope}

    def compose_for_agent(self, project_id: str, agent_id: str) -> str:
        return self.preview(project_id, agent_id=agent_id)["composed_prompt"]

    def get_effective_text(self, *, project_id: str, agent_id: str) -> str:
        state = self._load(project_id)
        skill = state["skills"].get(
            agent_id, _profile_record() | {"default_text": DEFAULT_AGENT_SKILLS.get(agent_id, "")},
        )
        return self._effective(skill)

    def get_inferred_text(self, project_id: str) -> str:
        return self._load(project_id)["project"].get("inferred_text", "")

    def get_manual_text(self, project_id: str) -> str:
        return self._load(project_id)["project"].get("manual_text", "")

    def set_inferred_text(self, project_id: str, text: str) -> dict[str, Any]:
        return self.update_profile(project_id, "project", text, source="inferred")

    async def review_commit_with_llm(
        self, *, project_id: str, interactions_text: str, current_profile: str,
        current_persona: str, manual_persona: str,
    ) -> CommitAdaptationReview:
        del project_id  # not yet part of the prompt; kept for a future per-project rubric override
        rubric = (
            _COMMIT_REVIEWER_SKILL_PATH.read_text(encoding="utf-8")
            if _COMMIT_REVIEWER_SKILL_PATH.exists() else
            "Review recent student interactions and propose memory/persona updates."
        )
        messages = [
            {"role": "system", "content": rubric},
            {
                "role": "user",
                "content": (
                    f"STUDENT-AUTHORED INTERACTIONS SINCE LAST REVIEW:\n{interactions_text}\n\n"
                    f"CURRENT DURABLE STUDENT PROFILE (Cognee):\n{current_profile}\n\n"
                    f"CURRENT INFERRED PERSONA TEXT:\n{current_persona or 'None yet.'}\n\n"
                    f"MANUAL PERSONA OVERRIDES (must not be contradicted):\n{manual_persona or 'None.'}\n\n"
                    "Propose memory candidates with exact interaction quotes as evidence, and rewrite the "
                    "complete inferred persona text (not a diff)."
                ),
            },
        ]
        client = self._client
        if client is None:
            from app.agents.cerebras_client import CerebrasClient

            client = CerebrasClient()
        return client.structured_complete(messages, CommitAdaptationReview)

    # evolve_from_commit() was retired: it was a non-LLM string template
    # (truncate the Evaluator's session_summary into the tutor skill slot)
    # superseded by review_commit_with_llm() + run_commit_adaptation(),
    # which actually reasons about accepted memory candidates rather than
    # blindly excerpting text. See
    # docs/commit-memory-adaptation-architecture.md.