exec-aide/src/conversation/tools.py

import json
import logging
import os
from datetime import datetime, timezone

import boto3

from shared.dynamo import (
    get_todays_messages, get_unanswered_threads, save_coordination,
    create_task, get_tasks, update_task_status, delete_task,
)
from shared.gmail import get_authenticated_service, fetch_thread
from shared.secrets import get_config, get_slack_token
from shared.slack import post_message

logger = logging.getLogger(__name__)

_lambda_client = None
_tool_context = {}


def set_tool_context(channel, thread_ts):
    global _tool_context
    _tool_context = {"channel": channel, "thread_ts": thread_ts}


TOOLS = [
    {
        "toolSpec": {
            "name": "get_inbox_summary",
            "description": "Get a summary of today's inbox: total emails processed, counts by classification (HIGH, NORMAL, LOW), and a list of the most recent high-priority items.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "date": {
                            "type": "string",
                            "description": "Date in YYYY-MM-DD format. Defaults to today.",
                        }
                    },
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "get_high_priority_emails",
            "description": "Get all HIGH-priority emails for a given date.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "date": {
                            "type": "string",
                            "description": "Date in YYYY-MM-DD format. Defaults to today.",
                        }
                    },
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "get_unanswered_threads",
            "description": "Get email threads where Adam hasn't replied in 24+ hours.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {},
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "get_bypassed_work_orders",
            "description": "Get emails sent directly to Adam without CC'ing the work-orders address. These should have gone through the work order system.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "date": {
                            "type": "string",
                            "description": "Date in YYYY-MM-DD format. Defaults to today.",
                        }
                    },
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "search_emails_by_sender",
            "description": "Search for emails from a specific sender by name or email address. Searches across all classified emails in the database.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "sender": {
                            "type": "string",
                            "description": "Sender name or email address to search for (case-insensitive partial match).",
                        },
                        "limit": {
                            "type": "integer",
                            "description": "Maximum number of results to return. Defaults to 10.",
                        },
                    },
                    "required": ["sender"],
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "get_email_thread_detail",
            "description": "Get the full message history of a specific email thread from Gmail, including all participants and message snippets.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "thread_id": {
                            "type": "string",
                            "description": "Gmail thread ID.",
                        }
                    },
                    "required": ["thread_id"],
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "trigger_daily_digest",
            "description": "Send Adam his daily inbox digest right now, regardless of the scheduled time.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {},
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "search_inbox",
            "description": "Search Adam's Gmail inbox with a query. Supports Gmail search syntax: from:, to:, subject:, after:, before:, is:unread, label:, has:attachment, etc.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "Gmail search query (e.g., 'from:jane after:2024/01/01 subject:invoice').",
                        },
                        "max_results": {
                            "type": "integer",
                            "description": "Maximum results to return. Defaults to 10.",
                        },
                    },
                    "required": ["query"],
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "coordinate_with_user",
            "description": "Start a coordination thread to interact with someone in the channel. Posts a message in the channel and starts a threaded conversation with the mentioned person. Use when Adam asks you to coordinate, schedule, or communicate with someone.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "flat_message": {
                            "type": "string",
                            "description": "Brief message to post flat in the channel (e.g., 'Coordinating a meeting with Jane').",
                        },
                        "thread_message": {
                            "type": "string",
                            "description": "The message to post in the thread, @mentioning the person (use Slack format <@USER_ID>).",
                        },
                    },
                    "required": ["flat_message", "thread_message"],
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "create_task",
            "description": "Create a new task/to-do item for Adam.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "title": {
                            "type": "string",
                            "description": "Task description.",
                        },
                        "due_date": {
                            "type": "string",
                            "description": "Optional due date in YYYY-MM-DD format.",
                        },
                    },
                    "required": ["title"],
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "list_tasks",
            "description": "List Adam's tasks. Can filter by status (open or completed).",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "status": {
                            "type": "string",
                            "description": "Filter by status: 'open' or 'completed'. Omit for all tasks.",
                        },
                    },
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "complete_task",
            "description": "Mark a task as completed.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "task_id": {
                            "type": "string",
                            "description": "The task ID (ULID) to mark as completed.",
                        },
                    },
                    "required": ["task_id"],
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "delete_task",
            "description": "Delete a task permanently.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "task_id": {
                            "type": "string",
                            "description": "The task ID (ULID) to delete.",
                        },
                    },
                    "required": ["task_id"],
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "create_reminder",
            "description": "Set a reminder that will send Adam a Slack DM at a specific time.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "text": {
                            "type": "string",
                            "description": "The reminder message text.",
                        },
                        "remind_at": {
                            "type": "string",
                            "description": "When to send the reminder in ISO 8601 format (e.g., 2026-05-05T09:00:00). Timezone is America/New_York.",
                        },
                    },
                    "required": ["text", "remind_at"],
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "get_calendar_events",
            "description": "Get Adam's upcoming calendar events for a date range.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "start_date": {
                            "type": "string",
                            "description": "Start date in YYYY-MM-DD format. Defaults to today.",
                        },
                        "end_date": {
                            "type": "string",
                            "description": "End date in YYYY-MM-DD format. Defaults to start_date + 1 day.",
                        },
                        "max_results": {
                            "type": "integer",
                            "description": "Maximum events to return. Defaults to 10.",
                        },
                    },
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "create_calendar_event",
            "description": "Create a new event on Adam's Google Calendar. Always confirm details with Adam before calling this.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "summary": {
                            "type": "string",
                            "description": "Event title.",
                        },
                        "start_time": {
                            "type": "string",
                            "description": "Start time in ISO 8601 (e.g., 2026-05-05T10:00:00).",
                        },
                        "end_time": {
                            "type": "string",
                            "description": "End time in ISO 8601 (e.g., 2026-05-05T11:00:00).",
                        },
                        "attendees": {
                            "type": "array",
                            "items": {"type": "string"},
                            "description": "Email addresses of attendees.",
                        },
                        "description": {
                            "type": "string",
                            "description": "Event description or notes.",
                        },
                    },
                    "required": ["summary", "start_time", "end_time"],
                }
            },
        }
    },
    {
        "toolSpec": {
            "name": "check_availability",
            "description": "Check Adam's calendar availability (free/busy) for a time range.",
            "inputSchema": {
                "json": {
                    "type": "object",
                    "properties": {
                        "start_time": {
                            "type": "string",
                            "description": "Range start in ISO 8601 format.",
                        },
                        "end_time": {
                            "type": "string",
                            "description": "Range end in ISO 8601 format.",
                        },
                    },
                    "required": ["start_time", "end_time"],
                }
            },
        }
    },
]


def _get_lambda_client():
    global _lambda_client
    if _lambda_client is None:
        _lambda_client = boto3.client("lambda")
    return _lambda_client


def _today():
    return datetime.now(timezone.utc).strftime("%Y-%m-%d")


def _format_email(msg):
    return {
        "message_id": msg.get("message_id") or msg.get("pk", "").replace("MSG#", ""),
        "thread_id": msg.get("thread_id", ""),
        "from": msg.get("from_name") or msg.get("from_email", ""),
        "from_email": msg.get("from_email", ""),
        "subject": msg.get("subject", ""),
        "snippet": msg.get("snippet", ""),
        "classification": msg.get("classification", ""),
        "classification_reason": msg.get("classification_reason", ""),
        "internal_date": int(msg.get("internal_date", 0)),
    }


def dispatch_tool(tool_name, tool_input):
    handlers = {
        "get_inbox_summary": _handle_inbox_summary,
        "get_high_priority_emails": _handle_high_priority,
        "get_unanswered_threads": _handle_unanswered,
        "get_bypassed_work_orders": _handle_bypassed,
        "search_emails_by_sender": _handle_search_sender,
        "get_email_thread_detail": _handle_thread_detail,
        "trigger_daily_digest": _handle_trigger_digest,
        "search_inbox": _handle_search_inbox,
        "coordinate_with_user": _handle_coordinate,
        "create_task": _handle_create_task,
        "list_tasks": _handle_list_tasks,
        "complete_task": _handle_complete_task,
        "delete_task": _handle_delete_task,
        "create_reminder": _handle_create_reminder,
        "get_calendar_events": _handle_get_calendar_events,
        "create_calendar_event": _handle_create_calendar_event,
        "check_availability": _handle_check_availability,
    }
    handler = handlers.get(tool_name)
    if not handler:
        return {"error": f"Unknown tool: {tool_name}"}
    try:
        return handler(tool_input)
    except Exception:
        logger.exception("Tool %s failed", tool_name)
        return {"error": f"Tool {tool_name} failed"}


def _handle_inbox_summary(inputs):
    date = inputs.get("date") or _today()
    messages = get_todays_messages(date)
    by_class = {}
    for msg in messages:
        c = msg.get("classification", "UNKNOWN")
        by_class[c] = by_class.get(c, 0) + 1

    high = [_format_email(m) for m in messages if m.get("classification") == "HIGH"]
    return {
        "date": date,
        "total": len(messages),
        "by_classification": by_class,
        "high_priority_emails": high[:10],
    }


def _handle_high_priority(inputs):
    date = inputs.get("date") or _today()
    messages = get_todays_messages(date)
    high = [_format_email(m) for m in messages if m.get("classification") == "HIGH"]
    return {"date": date, "count": len(high), "emails": high}


def _handle_unanswered(inputs):
    config = get_config()
    threads = get_unanswered_threads(config["unanswered_threshold_hours"])
    return {
        "count": len(threads),
        "threads": [
            {
                "thread_id": t.get("pk", "").replace("THD#", ""),
                "subject": t.get("subject", ""),
                "last_message_from": t.get("last_message_from", ""),
                "unanswered_since": t.get("unanswered_since", ""),
            }
            for t in threads
        ],
    }


def _handle_bypassed(inputs):
    date = inputs.get("date") or _today()
    messages = get_todays_messages(date)
    bypassed = [_format_email(m) for m in messages if m.get("bypassed_wo")]
    return {"date": date, "count": len(bypassed), "emails": bypassed}


def _handle_search_sender(inputs):
    from shared.dynamo import _get_table

    sender = inputs["sender"].lower()
    limit = inputs.get("limit", 10)

    results = []
    scan_kwargs = {
        "FilterExpression": "sk = :sk AND (contains(from_email, :sender) OR contains(from_name, :sender))",
        "ExpressionAttributeValues": {":sk": "MSG", ":sender": sender},
        "Limit": 500,
    }
    resp = _get_table().scan(**scan_kwargs)
    results.extend(resp.get("Items", []))
    while "LastEvaluatedKey" in resp and len(results) < limit:
        scan_kwargs["ExclusiveStartKey"] = resp["LastEvaluatedKey"]
        resp = _get_table().scan(**scan_kwargs)
        results.extend(resp.get("Items", []))

    results.sort(key=lambda m: int(m.get("internal_date", 0)), reverse=True)
    return {
        "sender_query": inputs["sender"],
        "count": len(results[:limit]),
        "emails": [_format_email(m) for m in results[:limit]],
    }


def _handle_thread_detail(inputs):
    service = get_authenticated_service()
    messages = fetch_thread(service, inputs["thread_id"])
    return {
        "thread_id": inputs["thread_id"],
        "message_count": len(messages),
        "messages": [
            {
                "from": m.get("from_name") or m.get("from_email", ""),
                "from_email": m.get("from_email", ""),
                "subject": m.get("subject", ""),
                "snippet": m.get("snippet", ""),
                "internal_date": m.get("internal_date", 0),
            }
            for m in messages
        ],
    }


def _handle_trigger_digest(inputs):
    resp = _get_lambda_client().invoke(
        FunctionName="exec-aide-daily-digest",
        InvocationType="RequestResponse",
    )
    payload = json.loads(resp["Payload"].read())
    return {
        "status": "sent",
        "high_priority": payload.get("high", 0),
        "bypassed_work_orders": payload.get("bypassed", 0),
        "unanswered_threads": payload.get("unanswered", 0),
        "normal": payload.get("normal", 0),
        "low": payload.get("low", 0),
    }


def _handle_search_inbox(inputs):
    service = get_authenticated_service()
    query = inputs["query"]
    max_results = inputs.get("max_results", 10)

    resp = service.users().messages().list(
        userId="me", q=query, maxResults=max_results
    ).execute()

    messages = []
    for msg_meta in resp.get("messages", []):
        msg = service.users().messages().get(
            userId="me", id=msg_meta["id"], format="metadata",
            metadataHeaders=["From", "To", "Subject", "Date"],
        ).execute()
        headers = {h["name"]: h["value"] for h in msg["payload"]["headers"]}
        messages.append({
            "id": msg["id"],
            "thread_id": msg["threadId"],
            "from": headers.get("From", ""),
            "to": headers.get("To", ""),
            "subject": headers.get("Subject", ""),
            "date": headers.get("Date", ""),
            "snippet": msg.get("snippet", ""),
        })

    return {"query": query, "count": len(messages), "messages": messages}


def _handle_coordinate(inputs):
    ctx = _tool_context
    if not ctx.get("channel"):
        return {"error": "Coordination is only available in channels, not DMs."}

    token = get_slack_token()
    channel = ctx["channel"]

    coord_ts = post_message(channel, inputs["flat_message"], token)
    post_message(channel, inputs["thread_message"], token, thread_ts=coord_ts)

    save_coordination(
        channel=channel,
        coord_thread_ts=coord_ts,
        conversation_thread_ts=ctx.get("thread_ts", ""),
        started_by="lauren",
    )

    return {
        "status": "coordinating",
        "message": "Thread started. I'll follow up when they respond.",
    }


# ── Task & Reminder Handlers ──────────────────────────────────


def _handle_create_task(inputs):
    task_id = create_task(inputs["title"], inputs.get("due_date"))
    return {"task_id": task_id, "status": "created", "title": inputs["title"]}


def _handle_list_tasks(inputs):
    status_filter = inputs.get("status")
    tasks = get_tasks(status_filter)
    return {
        "count": len(tasks),
        "tasks": [
            {
                "task_id": t["pk"].replace("TASK#", ""),
                "title": t["title"],
                "due_date": t.get("due_date", ""),
                "status": t["status"],
                "created_at": t.get("created_at", ""),
            }
            for t in tasks
        ],
    }


def _handle_complete_task(inputs):
    update_task_status(inputs["task_id"], "completed")
    return {"task_id": inputs["task_id"], "status": "completed"}


def _handle_delete_task(inputs):
    delete_task(inputs["task_id"])
    return {"task_id": inputs["task_id"], "status": "deleted"}


def _handle_create_reminder(inputs):
    import ulid as _ulid

    scheduler = boto3.client("scheduler")
    schedule_name = f"exec-aide-reminder-{_ulid.new()}"
    remind_at = inputs["remind_at"]

    scheduler.create_schedule(
        Name=schedule_name,
        ScheduleExpression=f"at({remind_at})",
        ScheduleExpressionTimezone="America/New_York",
        FlexibleTimeWindow={"Mode": "OFF"},
        Target={
            "Arn": os.environ["REMINDER_FN_ARN"],
            "RoleArn": os.environ["REMINDER_SCHEDULER_ROLE_ARN"],
            "Input": json.dumps({"text": inputs["text"]}),
        },
        ActionAfterCompletion="DELETE",
    )
    return {"status": "scheduled", "remind_at": remind_at, "text": inputs["text"]}


# ── Calendar Handlers ──────────────────────────────────────────


def _handle_get_calendar_events(inputs):
    from datetime import timedelta
    from shared.calendar import get_calendar_service, list_events

    start_date = inputs.get("start_date") or _today()
    end_date = inputs.get("end_date")
    if not end_date:
        d = datetime.strptime(start_date, "%Y-%m-%d")
        end_date = (d + timedelta(days=1)).strftime("%Y-%m-%d")

    time_min = f"{start_date}T00:00:00-04:00"
    time_max = f"{end_date}T23:59:59-04:00"
    max_results = inputs.get("max_results", 10)

    service = get_calendar_service()
    events = list_events(service, time_min, time_max, max_results)

    return {
        "count": len(events),
        "events": [
            {
                "id": e.get("id", ""),
                "summary": e.get("summary", "(no title)"),
                "start": e.get("start", {}).get("dateTime") or e.get("start", {}).get("date", ""),
                "end": e.get("end", {}).get("dateTime") or e.get("end", {}).get("date", ""),
                "attendees": [a.get("email", "") for a in e.get("attendees", [])],
                "location": e.get("location", ""),
                "status": e.get("status", ""),
            }
            for e in events
        ],
    }


def _handle_create_calendar_event(inputs):
    from shared.calendar import get_calendar_service, create_event

    service = get_calendar_service()
    event = create_event(
        service,
        summary=inputs["summary"],
        start={"dateTime": inputs["start_time"], "timeZone": "America/New_York"},
        end={"dateTime": inputs["end_time"], "timeZone": "America/New_York"},
        attendees=inputs.get("attendees"),
        description=inputs.get("description"),
    )
    return {
        "status": "created",
        "event_id": event["id"],
        "summary": event.get("summary", ""),
        "html_link": event.get("htmlLink", ""),
    }


def _handle_check_availability(inputs):
    from shared.calendar import get_calendar_service, check_availability

    service = get_calendar_service()
    busy = check_availability(service, inputs["start_time"], inputs["end_time"])
    return {
        "busy_slots": busy,
        "is_free": len(busy) == 0,
    }