pigway-frame-flow/pigway_frame_flow/providers/ical.py

"""Bounded read-only iCalendar expansion for the next seven days."""

import re
from datetime import UTC, datetime, timedelta
from zoneinfo import ZoneInfo

from dateutil.rrule import rrulestr

from pigway_frame_flow.errors import ErrorCode as E
from pigway_frame_flow.errors import FlowError


def calendar_events(raw, now=None):
    now = now or datetime.now(UTC)
    end = now + timedelta(days=7)
    text = re.sub(r"\r?\n[ \t]", "", raw.decode("utf-8-sig"))
    blocks = text.split("BEGIN:VEVENT")[1:]
    if len(blocks) > 2000:
        raise FlowError(E.INVALID_REQUEST, "Calendar exceeds 2000 events")

    def stamp(key, value):
        zone = re.search(r"TZID=([^;:]+)", key)
        tz = ZoneInfo(zone[1]) if zone else ZoneInfo("Asia/Shanghai")
        if len(value) == 8:
            return datetime.strptime(value, "%Y%m%d").replace(tzinfo=tz)
        fmt = "%Y%m%dT%H%M%SZ" if value.endswith("Z") else "%Y%m%dT%H%M%S"
        return datetime.strptime(value, fmt).replace(
            tzinfo=UTC if value.endswith("Z") else tz
        )

    parsed: list[tuple] = []
    for block in blocks:
        props: dict[str, list[tuple[str, str]]] = {}
        for line in block.split("END:VEVENT")[0].splitlines():
            if ":" in line:
                key, value = line.split(":", 1)
                props.setdefault(key.split(";")[0], []).append((key, value))
        if "DTSTART" not in props or "SUMMARY" not in props:
            continue
        start = stamp(*props["DTSTART"][0])
        uid = props.get("UID", [(None, str(len(parsed)))])[0][1]
        override = (
            stamp(*props["RECURRENCE-ID"][0]) if "RECURRENCE-ID" in props else None
        )
        parsed.append((props, start, uid, override))
    overrides = {
        (uid, original.astimezone(UTC)) for _, _, uid, original in parsed if original
    }
    result = []
    for props, start, uid, override in parsed:
        if props.get("STATUS", [(None, "")])[0][1] == "CANCELLED":
            continue
        if "EXRULE" in props:
            raise FlowError(
                E.UNSUPPORTED_CAPABILITY,
                "Use EXDATE instead of unsupported EXRULE",
                422,
            )
        occurrences = [start]
        if "RRULE" in props and not override:
            rule = props["RRULE"][0][1]
            if not re.search(r"(?:^|;)FREQ=(DAILY|WEEKLY|MONTHLY|YEARLY)(?:;|$)", rule):
                raise FlowError(
                    E.UNSUPPORTED_CAPABILITY,
                    "Calendar frequency must be daily or slower",
                    422,
                )
            if now.year - start.year > 150:
                raise FlowError(
                    E.INVALID_REQUEST, "Calendar recurrence starts too far in the past"
                )
            if len(rule) > 1000:
                raise FlowError(
                    E.INVALID_REQUEST, "Calendar recurrence rule is too long"
                )
            occurrences = []
            for occurrence in rrulestr(rule, dtstart=start).xafter(
                now - timedelta(days=1), count=1001, inc=True
            ):
                if occurrence >= end:
                    break
                occurrences.append(occurrence)
            if len(occurrences) > 1000:
                raise FlowError(
                    E.INVALID_REQUEST, "Calendar recurrence exceeds 1000 instances"
                )
        for key, value in props.get("RDATE", []):
            occurrences.extend(stamp(key, item) for item in value.split(","))
        excluded = {
            stamp(key, item).astimezone(UTC)
            for key, value in props.get("EXDATE", [])
            for item in value.split(",")
        }
        title = (
            props["SUMMARY"][0][1]
            .replace("\\n", " ")
            .replace("\\,", ",")
            .replace("\\;", ";")
            .replace("\\\\", "\\")
        )
        all_day = len(props["DTSTART"][0][1]) == 8
        for occurrence in sorted(set(occurrences)):
            instant = occurrence.astimezone(UTC)
            lower = (
                now.astimezone(occurrence.tzinfo).replace(
                    hour=0, minute=0, second=0, microsecond=0
                )
                if all_day
                else now
            )
            if instant in excluded or (not override and (uid, instant) in overrides):
                continue
            if lower <= instant < end:
                result.append({"title": title, "start": occurrence.isoformat()})
                if len(result) > 1000:
                    raise FlowError(
                        E.INVALID_REQUEST, "Calendar expansion exceeds 1000 instances"
                    )
    return {"events": sorted(result, key=lambda x: x["start"])}