afterhours-shift-manager/src/schedule.py

"""DynamoDB-backed shift schedule management.

Single-table design:
  PK=ROSTER   SK=<extension>   — employee roster entry
  PK=WEEKLY   SK=<DayName>     — default weekly schedule
  PK=OVERRIDE SK=<YYYY-MM-DD>  — date-specific override (or OPEN)
  PK=CONFIG   SK=CONFIG         — global settings
"""

import os

import boto3
from boto3.dynamodb.conditions import Key


FALLBACK_EXTENSION = "100"


class ShiftSchedule:
    def __init__(self, table_name: str | None = None):
        dynamodb = boto3.resource("dynamodb")
        self.table = dynamodb.Table(table_name or os.environ["SHIFT_TABLE"])

    # ── Roster ──────────────────────────────────────────────────────────

    def get_roster(self) -> list[dict]:
        resp = self.table.query(KeyConditionExpression=Key("PK").eq("ROSTER"))
        return resp["Items"]

    def get_employee_by_extension(self, extension: str) -> dict | None:
        resp = self.table.get_item(Key={"PK": "ROSTER", "SK": extension})
        return resp.get("Item")

    def get_employee_by_slack_id(self, slack_user_id: str) -> dict | None:
        for item in self.get_roster():
            if item.get("slack_user_id") == slack_user_id:
                return item
        return None

    def register_user(self, slack_user_id: str, extension: str) -> dict | None:
        employee = self.get_employee_by_extension(extension)
        if not employee:
            return None
        self.table.update_item(
            Key={"PK": "ROSTER", "SK": extension},
            UpdateExpression="SET slack_user_id = :sid",
            ExpressionAttributeValues={":sid": slack_user_id},
        )
        employee["slack_user_id"] = slack_user_id
        return employee

    # ── Weekly schedule ─────────────────────────────────────────────────

    def get_weekly_schedule(self) -> dict[str, dict]:
        resp = self.table.query(KeyConditionExpression=Key("PK").eq("WEEKLY"))
        return {item["SK"]: item for item in resp["Items"]}

    # ── Overrides ───────────────────────────────────────────────────────

    def get_override(self, date_str: str) -> dict | None:
        resp = self.table.get_item(Key={"PK": "OVERRIDE", "SK": date_str})
        return resp.get("Item")

    def set_override(self, date_str: str, extension: str, name: str) -> None:
        self.table.put_item(
            Item={
                "PK": "OVERRIDE",
                "SK": date_str,
                "extension": extension,
                "name": name,
            }
        )

    def mark_open(self, date_str: str) -> None:
        self.table.put_item(
            Item={
                "PK": "OVERRIDE",
                "SK": date_str,
                "extension": "OPEN",
                "name": "Open",
            }
        )

    def remove_override(self, date_str: str) -> None:
        self.table.delete_item(Key={"PK": "OVERRIDE", "SK": date_str})

    # ── Resolution ──────────────────────────────────────────────────────

    def resolve_shift(self, date_str: str, day_name: str) -> tuple[str, str, str]:
        """Return (extension, name, source) for a given date.

        source is one of: 'override', 'available', 'weekly'.
        'available' means no one is assigned — the fallback extension handles
        it silently on the 3CX side, but users just see it as open for pickup.
        """
        override = self.get_override(date_str)
        if override:
            if override["extension"] == "OPEN":
                return (FALLBACK_EXTENSION, "Available", "available")
            return (override["extension"], override["name"], "override")

        weekly = self.get_weekly_schedule()
        if day_name in weekly:
            item = weekly[day_name]
            if item["extension"] == FALLBACK_EXTENSION:
                return (FALLBACK_EXTENSION, "Available", "available")
            return (item["extension"], item.get("name", "Unknown"), "weekly")

        return (FALLBACK_EXTENSION, "Available", "available")

    # ── Config ──────────────────────────────────────────────────────────

    def get_config(self) -> dict:
        resp = self.table.get_item(Key={"PK": "CONFIG", "SK": "CONFIG"})
        return resp.get("Item", {})