afterhours-shift-manager/src/shared/python/shared/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, shift_type: str = "night") -> dict | None:
        sk = f"{date_str}-DAY" if shift_type == "day" else date_str
        resp = self.table.get_item(Key={"PK": "OVERRIDE", "SK": sk})
        return resp.get("Item")

    def set_override(
        self, date_str: str, extension: str, name: str, shift_type: str = "night"
    ) -> None:
        sk = f"{date_str}-DAY" if shift_type == "day" else date_str
        self.table.put_item(
            Item={
                "PK": "OVERRIDE",
                "SK": sk,
                "extension": extension,
                "name": name,
            }
        )

    def claim_open_shift(
        self, date_str: str, extension: str, name: str, shift_type: str = "night"
    ) -> bool:
        """Atomically claim a shift only if it is currently open.

        Succeeds when no override exists (weekly fallback) or the override
        is explicitly OPEN. Fails if someone else already claimed it.
        """
        sk = f"{date_str}-DAY" if shift_type == "day" else date_str
        try:
            self.table.put_item(
                Item={
                    "PK": "OVERRIDE",
                    "SK": sk,
                    "extension": extension,
                    "name": name,
                },
                ConditionExpression="attribute_not_exists(PK) OR extension = :open",
                ExpressionAttributeValues={":open": "OPEN"},
            )
            return True
        except self.table.meta.client.exceptions.ConditionalCheckFailedException:
            return False

    def mark_open(self, date_str: str, shift_type: str = "night") -> None:
        sk = f"{date_str}-DAY" if shift_type == "day" else date_str
        self.table.put_item(
            Item={
                "PK": "OVERRIDE",
                "SK": sk,
                "extension": "OPEN",
                "name": "Open",
            }
        )

    def remove_override(self, date_str: str, shift_type: str = "night") -> None:
        sk = f"{date_str}-DAY" if shift_type == "day" else date_str
        self.table.delete_item(Key={"PK": "OVERRIDE", "SK": sk})

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

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

        shift_type: 'night' (after-hours) or 'day' (weekend 8am-5pm)
        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, shift_type)
        if override:
            if override["extension"] == "OPEN":
                return (FALLBACK_EXTENSION, "Available", "available")
            return (override["extension"], override["name"], "override")

        weekly = self.get_weekly_schedule()
        weekly_key = f"{day_name}-Day" if shift_type == "day" else day_name
        if weekly_key in weekly:
            item = weekly[weekly_key]
            if item["extension"] == FALLBACK_EXTENSION:
                return (FALLBACK_EXTENSION, "Available", "available")
            return (item["extension"], item.get("name", "Unknown"), "weekly")

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

    # ── Schedule post tracking ───────────────────────────────────────────

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

    def save_schedule_post(
        self, channel_id: str, message_ts: str, week_start: str
    ) -> None:
        self.table.put_item(
            Item={
                "PK": "SCHEDULE_POST",
                "SK": channel_id,
                "message_ts": message_ts,
                "week_start": week_start,
            }
        )

    # ── Pay records ─────────────────────────────────────────────────────

    def get_pay_record(self, week_key: str) -> dict | None:
        """Get a pay record by week key (e.g. '2026-04-06')."""
        resp = self.table.get_item(Key={"PK": "PAY", "SK": week_key})
        return resp.get("Item")

    def save_pay_record(self, week_key: str, record: dict) -> None:
        """Save a weekly pay summary. week_key is the Monday date string."""
        self.table.put_item(Item={"PK": "PAY", "SK": week_key, **record})

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

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

    def get_admin_users(self) -> list[str]:
        config = self.get_config()
        return config.get("admin_users", [])

    def add_roster_entry(self, extension: str, name: str) -> None:
        self.table.put_item(
            Item={
                "PK": "ROSTER",
                "SK": extension,
                "name": name,
                "extension": extension,
                "slack_user_id": "",
            }
        )

    def remove_roster_entry(self, extension: str) -> None:
        self.table.delete_item(Key={"PK": "ROSTER", "SK": extension})

    def rename_roster_entry(self, extension: str, new_name: str) -> None:
        self.table.update_item(
            Key={"PK": "ROSTER", "SK": extension},
            UpdateExpression="SET #n = :name",
            ExpressionAttributeNames={"#n": "name"},
            ExpressionAttributeValues={":name": new_name},
        )

    def set_default_shift_rate(self, rate: float) -> None:
        """Set the global default shift rate in config."""
        self.table.update_item(
            Key={"PK": "CONFIG", "SK": "CONFIG"},
            UpdateExpression="SET shift_rate = :r",
            ExpressionAttributeValues={":r": str(rate)},
        )

    def set_employee_shift_rate(self, extension: str, rate: float) -> None:
        """Set a per-person shift rate on a roster entry."""
        self.table.update_item(
            Key={"PK": "ROSTER", "SK": extension},
            UpdateExpression="SET shift_rate = :r",
            ExpressionAttributeValues={":r": str(rate)},
        )

    def get_shift_rate(self, extension: str | None = None) -> float:
        """Get the per-shift pay rate.

        Checks the roster entry for a person-specific rate first,
        then falls back to the global config rate.
        """
        if extension:
            employee = self.get_employee_by_extension(extension)
            if employee and employee.get("shift_rate"):
                return float(employee["shift_rate"])
        config = self.get_config()
        return float(config.get("shift_rate", 0))