shiftcreator.py

import asyncio
import os
from datetime import datetime, timedelta, time
from dotenv import load_dotenv
import pytz

from rich.console import Console
from rich.panel import Panel
from rich.prompt import Prompt, IntPrompt, Confirm
from rich.table import Table
from rich.progress import Progress

from msgraph import GraphServiceClient
from azure.identity import InteractiveBrowserCredential
from msgraph.generated.models.event import Event
from msgraph.generated.models.date_time_time_zone import DateTimeTimeZone
from msgraph.generated.models.item_body import ItemBody
from msgraph.generated.models.body_type import BodyType

console = Console()

# -----------------------------
# DEFAULT CONFIG
# -----------------------------

# Load .env file
load_dotenv()

CLIENT_ID = os.getenv("CLIENT_ID")
TENANT_ID = os.getenv("TENANT_ID")
                      
SHIFT_PATTERN = [
    ('night', 3),
    ('day', 2),
    ('night', 4),
    ('day', 2),
    ('day', 3)
]

SCOPES = ["Calendars.ReadWrite"]


# -----------------------------
# UTILITY FUNCTIONS
# -----------------------------

# Get calendar id by name
async def get_calendar_id(client, name):
    """ Finds the calendar id of the provided calendar name using Microsoft Graph

    Args:
        client (GraphServiceClient): Microsoft Graph Client
        name (str): Calendar name

    Raises:
        Exception: Can't find a calendar with the provided name

    Returns:
        int: Calendar id
    """

    # Get a list of all the users calendars
    calendars = await client.me.calendars.get()

    # Loop through the calendars and see if we can find one which matches the name provided
    for cal in calendars.value:
        if cal.name.lower() == name.lower():
            return cal.id
    raise Exception(f"Calendar '{name}' not found")

# Create shifts based on pattern and provided parameters 
def generate_shifts(start_date, days_to_cover, day_start, night_start, duration, timezone):
    """Create shifts based on pattern and provided parameters 

    Args:
        start_date (Date): Date to start generating from (this should be the first shift)
        days_to_cover (int): Days to create shifts for 
        day_start (int): Day shift start time
        night_start (int): Night shift start time
        duration (int): Shift duration 
        timezone (str): Timezone to use

    Returns:
        list: Shifts 
    """
    shifts = []
    current_date = start_date
    total_days = 0
    pattern_index = 0
    tz = pytz.timezone(timezone)

    while total_days < days_to_cover:
        shift_type, length = SHIFT_PATTERN[pattern_index]

        for _ in range(length):
            if total_days >= days_to_cover:
                break

            if shift_type == 'night':
                start_naive = datetime.combine(current_date, time(hour=night_start))
            else:
                start_naive = datetime.combine(current_date, time(hour=day_start))

            shift_start = tz.localize(start_naive)
            shift_end = shift_start + timedelta(hours=duration)

            shifts.append({
                "subject": f"{shift_type.capitalize()} Shift",
                "start": shift_start,
                "end": shift_end,
                "category": "Night Shift" if shift_type == "night" else "Day Shift"
            })

            current_date += timedelta(days=1)
            total_days += 1

        current_date += timedelta(days=length)
        total_days += length
        pattern_index = (pattern_index + 1) % len(SHIFT_PATTERN)

    return shifts

# Create calendar events
async def create_event(client, calendar_id, shift, timezone, reminder, dry_run=False):
    """_summary_

    Args:
        client (GraphServiceClient): Microsoft Graph Client
        calendar_id (int): Calendar ID
        shift (dict): Shift details to use
        timezone (string): Timezone to use
        reminder (int): Time in minutes to remind before
        dry_run (bool, optional): _description_. Defaults to False.
    """
    if dry_run:
        return

    request_body = Event(
        subject=shift["subject"],
        start=DateTimeTimeZone(
            date_time=shift["start"].isoformat(),
            time_zone=timezone
        ),
        end=DateTimeTimeZone(
            date_time=shift["end"].isoformat(),
            time_zone=timezone
        ),
        categories=[shift["category"]],
        reminder_minutes_before_start=reminder
    )

    await client.me.calendars.by_calendar_id(calendar_id).events.post(request_body)


# -----------------------------
# ASYNC MAIN APPLICATION
# -----------------------------

async def main():
    """Application main
    """

    console.print("[bold cyan]ShiftCreator[/bold cyan] — generate Outlook events for your shift pattern\n")

    console.print(
        "[bold yellow]Prerequisites:[/bold yellow]\n"
        "• Outlook calendar must exist (default: 'work')\n"
        "• Categories 'Day Shift' and 'Night Shift' must already exist\n"
        "• You will sign in using Microsoft Interactive Login\n"
    )

    # Shift pattern table (smaller, no title)
    pattern_table = Table(show_header=True, header_style="bold cyan")
    pattern_table.add_column("Shift Type", style="cyan", no_wrap=True)
    pattern_table.add_column("Length (days)", style="green", no_wrap=True)

    for shift_type, length in SHIFT_PATTERN:
        pattern_table.add_row(shift_type.capitalize(), str(length))

    console.print(pattern_table)

    console.print(
        f"[bold]Important:[/bold] When asked for the shift pattern start date, "
        f"enter the date on which the first block begins — "
    )   

    # -----------------------------
    # MODE SELECTION
    # -----------------------------
    mode = Prompt.ask(
        "[bold cyan]Choose mode[/bold cyan]",
        choices=["simple", "advanced"],
        default="simple"
    )

    # -----------------------------
    # SIMPLE MODE
    # -----------------------------
    if mode == "simple":
        calendar_name = "work"
        day_start = 6       # 6 am
        night_start = 18    # 6 pm
        duration = 12       # 12 hours 
        reminder = 60       # 60 minutes / 1 hour before
        timezone_str = 'Europe/London'

        start_date_str = Prompt.ask("Enter start date [dim](DD-MM-YYYY)[/dim]")
        days = IntPrompt.ask("Enter number of days to generate", default=30)

    # -----------------------------
    # ADVANCED MODE
    # -----------------------------
    else:
        console.print("[cyan]Shift details...[/cyan]")
        day_start = IntPrompt.ask("Day shift start hour", default=6)
        night_start = IntPrompt.ask("Night shift start hour", default=18)
        duration = IntPrompt.ask("Shift duration (hours)", default=12)
        console.print("[cyan]Generation settings...[/cyan]")
        calendar_name = Prompt.ask("Calendar name [dim](case insensitive)[/dim]", default="work")
        start_date_str = Prompt.ask("Start date [dim](DD-MM-YYYY)[/dim]")
        days = IntPrompt.ask("Number of days to generate")
        timezone_str = Prompt.ask("Timezone", default='Europe/London')
        reminder = IntPrompt.ask("Time in minutes to remind before event")


    # Validate date
    try:
        start_date = datetime.strptime(start_date_str, "%d-%m-%Y").date()
    except ValueError:
        console.print("[red]Invalid date format. Use DD-MM-YYYY.[/red]")
        return
    
    # Validate timezone
    if timezone_str in pytz.all_timezones:
        timezone = timezone_str
    else:
        console.print("[red]Invalid timezone. Please try again.[/red]")
        return

    # Authenticate
    console.print("[cyan]Opening browser for authentication...[/cyan]")
    credentials = InteractiveBrowserCredential(client_id=CLIENT_ID, tenant_id=TENANT_ID)
    client = GraphServiceClient(credentials=credentials, scopes=SCOPES)

    # Force authentication to complete
    await client.me.get()
    
    # Get calendar
    console.print("[cyan]Fetching calendar...[/cyan]")
    calendar_id = await get_calendar_id(client, calendar_name)
    console.print(f"[green]Found calendar:[/green] {calendar_name}")

    # Generate shifts
    shifts = generate_shifts(start_date, days, day_start, night_start, duration, timezone)

    # Preview table
    table = Table(title="Generated Shifts")
    table.add_column("Subject")
    table.add_column("Start")
    table.add_column("End")
    table.add_column("Category")

    for s in shifts:
        table.add_row(
            s["subject"],
            s["start"].strftime("%Y-%m-%d %H:%M"),
            s["end"].strftime("%Y-%m-%d %H:%M"),
            s["category"]
        )

    console.print(table)

    if not Confirm.ask("Proceed with creating these events"):
        console.print("[yellow]Cancelled.[/yellow]")
        return

    # Create events
    console.print("[cyan]Creating events...[/cyan]")

    with Progress() as progress:
        task = progress.add_task("Creating...", total=len(shifts))
        for shift in shifts:
            await create_event(client, calendar_id, shift, timezone, reminder)
            progress.update(task, advance=1)

    console.print(Panel(
        "[bold green]All events created successfully![/bold green]",
        border_style="green"
    ))


# -----------------------------
# ENTRY POINT
# -----------------------------

if __name__ == "__main__":
    asyncio.run(main())