backend/mcp/tool_executor.py

"""Execution helpers for MCP-style tools."""

from __future__ import annotations

from typing import Callable, Dict

from ..tools.analysis_engine import AnalysisEngine
from ..tools.base import BlueprintResult
from .actions import Action, PlannerContext
from .context_resolver import resolve_relative_day


class ToolExecutor:
    """Dispatch actions emitted by the planner to the analysis engine."""

    def __init__(self, engine: AnalysisEngine) -> None:
        self.engine = engine
        self._handlers: Dict[str, Callable[[Action, PlannerContext], BlueprintResult]] = {
            # Core unified handlers
            "lookup_person": self._run_lookup_person,
            "lookup_location": self._run_location,
            "lookup_center": self._run_center,
            "lookup_advisorship": self._run_advisorship,
            "lookup_faculty_topic": self._run_faculty_topic,
            "lookup_staff_support": self._run_staff_support,
            # Office hours
            "lookup_office_hours": self._run_office_hours,
            # Events
            "list_events": self._run_list_events,
            # Noop
            "noop": self._run_noop,
        }

    def execute(self, action: Action, context: PlannerContext) -> BlueprintResult:
        """
        Execute an action and return the resulting blueprint output.

        Unknown actions fall back to a no-op response so the caller can
        surface a helpful error message instead of crashing.
        """
        handler = self._handlers.get(action.type, self._run_unknown)
        return handler(action, context)

    # Internal handlers ------------------------------------------------- #
    def _run_faculty_topic(self, action: Action, _: PlannerContext) -> BlueprintResult:
        topic = (action.arguments.get("topic") or "").strip()
        if not topic:
            return BlueprintResult(
                "faculty_by_topic",
                action.arguments,
                facts=[],
                notes=["Tell me which research area you'd like to explore."],
            )
        return self.engine.run("faculty_by_topic", topic=topic)

    def _run_location(self, action: Action, context: PlannerContext) -> BlueprintResult:
        name = (action.arguments.get("name") or action.arguments.get("person") or "").strip()
        if not name and action.arguments.get("use_last_subject"):
            if context.last_subject:
                name = context.last_subject
                action.arguments["name"] = name
        if not name:
            return BlueprintResult(
                "location",
                action.arguments,
                facts=[],
                notes=["Let me know who you want directions to."],
            )
        return self.engine.run("location", name=name)

    def _run_center(self, action: Action, context: PlannerContext) -> BlueprintResult:
        faculty = (action.arguments.get("faculty") or action.arguments.get("name") or "").strip()
        center = (action.arguments.get("center") or "").strip()
        
        if not faculty and not center and action.arguments.get("use_last_subject"):
            if context.last_subject:
                faculty = context.last_subject
                action.arguments["faculty"] = faculty
        
        if not faculty and not center:
            return BlueprintResult(
                "center",
                action.arguments,
                facts=[],
                notes=["Tell me a faculty name or center name."],
            )
        
        params = {}
        if center:
            params["center"] = center
        if faculty:
            params["faculty"] = faculty
        return self.engine.run("center", **params)

    def _run_advisorship(self, action: Action, context: PlannerContext) -> BlueprintResult:
        name = (
            action.arguments.get("name") or 
            action.arguments.get("student") or 
            action.arguments.get("faculty") or 
            ""
        ).strip()
        if not name and action.arguments.get("use_last_subject"):
            if context.last_subject:
                name = context.last_subject
                action.arguments["name"] = name
        if not name:
            return BlueprintResult(
                "advisorship",
                action.arguments,
                facts=[],
                notes=["Tell me which student or faculty member you're asking about."],
            )
        return self.engine.run("advisorship", name=name)

    def _run_list_events(self, action: Action, _: PlannerContext) -> BlueprintResult:
        keyword = (action.arguments.get("keyword") or "").strip()
        if keyword:
            action.arguments["keyword"] = keyword
        else:
            action.arguments.pop("keyword", None)
        params = {}
        if keyword:
            params["keyword"] = keyword
        return self.engine.run("upcoming_events", **params)

    def _run_staff_support(self, action: Action, _: PlannerContext) -> BlueprintResult:
        topic = (action.arguments.get("topic") or action.arguments.get("need") or action.arguments.get("keyword") or "").strip()
        if not topic:
            return BlueprintResult(
                "staff_support",
                action.arguments,
                facts=[],
                notes=["Tell me what kind of help you need (finance, advising, travel, etc.)."],
            )
        action.arguments["topic"] = topic
        return self.engine.run("staff_support", topic=topic)

    def _run_office_hours(self, action: Action, context: PlannerContext) -> BlueprintResult:
        class_name = (action.arguments.get("class_name") or action.arguments.get("course") or "").strip()
        person = (action.arguments.get("person") or action.arguments.get("name") or "").strip()
        if not class_name and not person and action.arguments.get("use_last_subject"):
            if getattr(context, "last_class", None):
                class_name = context.last_class
                action.arguments["class_name"] = class_name
            elif context.last_subject:
                ls = context.last_subject
                if any(c.isdigit() for c in ls) or (ls and ls.upper().startswith("CS")):
                    class_name = ls
                    action.arguments["class_name"] = class_name
                else:
                    person = ls
                    action.arguments["person"] = person
        if not class_name and not person:
            return BlueprintResult(
                "office_hours",
                action.arguments,
                facts=[],
                notes=["Tell me which course or whose office hours you'd like."],
            )
        params: dict = {}
        if class_name:
            params["class_name"] = class_name
        if person:
            params["person"] = person
        day = (action.arguments.get("day") or "").strip()
        if day:
            resolved = resolve_relative_day(day)
            if resolved:
                day = resolved
                action.arguments["day"] = day
            params["day"] = day
        return self.engine.run("office_hours", **params)

    def _run_lookup_person(self, action: Action, context: PlannerContext) -> BlueprintResult:
        name = (action.arguments.get("name") or "").strip()
        if not name and action.arguments.get("use_last_subject"):
            if context.last_subject:
                name = context.last_subject
                action.arguments["name"] = name
        if not name:
            return BlueprintResult(
                "person_lookup",
                action.arguments,
                facts=[],
                notes=["Tell me which person you'd like me to look up (faculty, staff, or student)."],
            )
        return self.engine.run("person_lookup", name=name)

    def _run_noop(self, action: Action, _: PlannerContext) -> BlueprintResult:
        message = action.arguments.get("message") or "I'm not sure how to help with that yet."
        return BlueprintResult("noop", action.arguments, facts=[], notes=[message])

    def _run_unknown(self, action: Action, _: PlannerContext) -> BlueprintResult:
        return BlueprintResult(
            "noop",
            action.arguments,
            facts=[],
            notes=[f"I'm not set up to handle '{action.type}' yet."],
        )