breakpilot-lehrer/timetable-solver-service/app/repository.py

"""Read stammdaten + constraints from PostgreSQL and turn them into Timefold
domain objects. Used by runner.py to build a Timetable problem instance.

Ownership is enforced via created_by_user_id everywhere — the solver only
sees data belonging to the Rektor who triggered the solve.
"""

from typing import Any

import asyncpg

from .domain import Lesson, Room, SchoolClass, Subject, Teacher, Timeslot, Timetable
from .rules import (
    RoomRequiresTypeRule, RoomUnavailableRule,
    SubjectPreferredPeriodRule, TeacherExcludedRoomRule,
    TeacherUnavailableDayRule, TeacherUnavailableWindowRule,
)


async def build_problem(
    pool: asyncpg.Pool,
    user_id: str,
    parent_solution_id: str | None = None,
) -> Timetable:
    async with pool.acquire() as conn:
        timeslots = await _load_timeslots(conn, user_id)
        rooms = await _load_rooms(conn, user_id)
        teachers = await _load_teachers(conn, user_id)
        classes = await _load_classes(conn, user_id)
        subjects = await _load_subjects(conn, user_id)
        lessons = await _build_lessons(conn, user_id, classes, subjects, teachers)
        rules = await _load_rules(conn, user_id)

        if parent_solution_id:
            await _inherit_pinned_from_parent(
                conn, parent_solution_id, lessons, timeslots, rooms,
            )

    return Timetable(
        timeslots=timeslots,
        rooms=rooms,
        teachers=teachers,
        classes=classes,
        subjects=subjects,
        lessons=lessons,
        score=None,
        **rules,
    )


async def _inherit_pinned_from_parent(
    conn: asyncpg.Connection,
    parent_solution_id: str,
    lessons: list[Lesson],
    timeslots: list[Timeslot],
    rooms: list[Room],
) -> None:
    """Apply pinned lessons from a parent solution onto the new problem.

    Matching rule: parent pinned lesson maps to a new Lesson with the same
    (class_id, subject_id) that hasn't already received a pinned assignment.
    Greedy first-fit so the count matches the parent's pinned count up to
    the curriculum's weekly_hours per (class, subject). If curriculum
    changed between solves, surplus pinned rows are silently dropped.
    """
    rows = await conn.fetch("""
        SELECT class_id::text, subject_id::text, room_id::text,
               day_of_week, period_index
        FROM tt_lesson
        WHERE solution_id = $1 AND pinned = true
    """, parent_solution_id)

    if not rows:
        return

    ts_by_dp = {(t.day_of_week, t.period_index): t for t in timeslots}
    room_by_id = {r.id: r for r in rooms}
    used: set[str] = set()

    for r in rows:
        ts = ts_by_dp.get((r["day_of_week"], r["period_index"]))
        if ts is None:
            continue  # period was deleted in the meantime
        # Find a not-yet-pinned lesson with matching class+subject.
        candidate: Lesson | None = None
        for lesson in lessons:
            if lesson.id in used:
                continue
            if (lesson.school_class.id == r["class_id"] and
                    lesson.subject.id == r["subject_id"]):
                candidate = lesson
                break
        if candidate is None:
            continue
        candidate.timeslot = ts
        room_id = r["room_id"]
        if room_id:
            candidate.room = room_by_id.get(room_id)
        candidate.pinned = True
        used.add(candidate.id)


async def _load_timeslots(conn: asyncpg.Connection, user_id: str) -> list[Timeslot]:
    rows = await conn.fetch("""
        SELECT id::text, day_of_week, period_index,
               to_char(start_time, 'HH24:MI') AS st,
               to_char(end_time, 'HH24:MI') AS et,
               is_break
        FROM tt_period
        WHERE created_by_user_id = $1 AND is_break = false
        ORDER BY day_of_week, period_index
    """, user_id)
    return [
        Timeslot(id=r["id"], day_of_week=r["day_of_week"], period_index=r["period_index"],
                 start_time=r["st"], end_time=r["et"])
        for r in rows
    ]


async def _load_rooms(conn: asyncpg.Connection, user_id: str) -> list[Room]:
    rows = await conn.fetch("""
        SELECT id::text, name, COALESCE(room_type, '') AS rt
        FROM tt_room WHERE created_by_user_id = $1 ORDER BY name
    """, user_id)
    return [Room(id=r["id"], name=r["name"], room_type=r["rt"]) for r in rows]


async def _load_teachers(conn: asyncpg.Connection, user_id: str) -> list[Teacher]:
    rows = await conn.fetch("""
        SELECT id::text, first_name, last_name, short_code
        FROM tt_teacher WHERE created_by_user_id = $1 ORDER BY last_name, first_name
    """, user_id)
    return [Teacher(id=r["id"], first_name=r["first_name"], last_name=r["last_name"], short_code=r["short_code"]) for r in rows]


async def _load_classes(conn: asyncpg.Connection, user_id: str) -> list[SchoolClass]:
    rows = await conn.fetch("""
        SELECT id::text, name, grade_level
        FROM tt_class WHERE created_by_user_id = $1 ORDER BY grade_level, name
    """, user_id)
    return [SchoolClass(id=r["id"], name=r["name"], grade_level=r["grade_level"]) for r in rows]


async def _load_subjects(conn: asyncpg.Connection, user_id: str) -> list[Subject]:
    rows = await conn.fetch("""
        SELECT id::text, name, short_code, COALESCE(required_room_type, '') AS rt
        FROM tt_subject WHERE created_by_user_id = $1 ORDER BY name
    """, user_id)
    return [Subject(id=r["id"], name=r["name"], short_code=r["short_code"], required_room_type=r["rt"]) for r in rows]


async def _build_lessons(
    conn: asyncpg.Connection,
    user_id: str,
    classes: list[SchoolClass],
    subjects: list[Subject],
    teachers: list[Teacher],
) -> list[Lesson]:
    """Materialise curriculum × assignment into Lesson instances.

    For each (class, subject) row in tt_curriculum with weekly_hours=N, we
    create N Lesson rows. The teacher is the one assigned in tt_assignment
    for the same (class, subject) — there must be exactly one, else the
    lesson can't be scheduled and is skipped (the UI surfaces this gap).
    """
    rows = await conn.fetch("""
        SELECT cu.class_id::text, cu.subject_id::text, cu.weekly_hours,
               a.teacher_id::text
        FROM tt_curriculum cu
        JOIN tt_class cl ON cu.class_id = cl.id
        LEFT JOIN tt_assignment a
          ON a.class_id = cu.class_id AND a.subject_id = cu.subject_id
        WHERE cl.created_by_user_id = $1
    """, user_id)

    class_by_id = {c.id: c for c in classes}
    subject_by_id = {s.id: s for s in subjects}
    teacher_by_id = {t.id: t for t in teachers}

    lessons: list[Lesson] = []
    counter = 0
    for r in rows:
        cls = class_by_id.get(r["class_id"])
        sub = subject_by_id.get(r["subject_id"])
        tch = teacher_by_id.get(r["teacher_id"]) if r["teacher_id"] else None
        if cls is None or sub is None or tch is None:
            # Missing assignment — solver can't schedule without a teacher.
            continue
        for _ in range(int(r["weekly_hours"])):
            lessons.append(Lesson(
                id=f"L{counter}-{cls.id[:6]}-{sub.id[:6]}",
                school_class=cls,
                subject=sub,
                teacher=tch,
            ))
            counter += 1
    return lessons


async def _load_rules(conn: asyncpg.Connection, user_id: str) -> dict[str, list[Any]]:
    """Pull the subset of constraint tables the constraint provider uses."""
    rules: dict[str, list[Any]] = {}

    rows = await conn.fetch("""
        SELECT teacher_id::text, day_of_week, is_hard, weight
        FROM tt_constraint_teacher_unavailable_day
        WHERE created_by_user_id = $1 AND active = true
    """, user_id)
    rules["teacher_unavailable_days"] = [TeacherUnavailableDayRule(**dict(r)) for r in rows]

    rows = await conn.fetch("""
        SELECT teacher_id::text, day_of_week,
               to_char(start_time, 'HH24:MI') AS start_time,
               to_char(end_time, 'HH24:MI') AS end_time,
               is_hard, weight
        FROM tt_constraint_teacher_unavailable_window
        WHERE created_by_user_id = $1 AND active = true
    """, user_id)
    rules["teacher_unavailable_windows"] = [TeacherUnavailableWindowRule(**dict(r)) for r in rows]

    rows = await conn.fetch("""
        SELECT teacher_id::text, room_id::text, is_hard, weight
        FROM tt_constraint_teacher_excluded_room
        WHERE created_by_user_id = $1 AND active = true
    """, user_id)
    rules["teacher_excluded_rooms"] = [TeacherExcludedRoomRule(**dict(r)) for r in rows]

    rows = await conn.fetch("""
        SELECT room_id::text, day_of_week, period_index, is_hard, weight
        FROM tt_constraint_room_unavailable
        WHERE created_by_user_id = $1 AND active = true
    """, user_id)
    rules["room_unavailables"] = [RoomUnavailableRule(**dict(r)) for r in rows]

    rows = await conn.fetch("""
        SELECT subject_id::text, period_from, period_to, is_hard, weight
        FROM tt_constraint_subject_preferred_period
        WHERE created_by_user_id = $1 AND active = true
    """, user_id)
    rules["subject_preferred_periods"] = [SubjectPreferredPeriodRule(**dict(r)) for r in rows]

    rows = await conn.fetch("""
        SELECT subject_id::text, room_type, is_hard, weight
        FROM tt_constraint_room_requires_type
        WHERE created_by_user_id = $1 AND active = true
    """, user_id)
    rules["room_requires_types"] = [RoomRequiresTypeRule(**dict(r)) for r in rows]

    return rules


async def persist_solution(
    pool: asyncpg.Pool,
    solution_id: str,
    timetable: Timetable,
    hard_score: int,
    soft_score: int,
) -> None:
    """Write the solver result back to tt_solution + tt_lesson."""
    async with pool.acquire() as conn:
        async with conn.transaction():
            await conn.execute("""
                UPDATE tt_solution
                SET status = 'completed',
                    hard_score = $2,
                    soft_score = $3,
                    finished_at = NOW()
                WHERE id = $1
            """, solution_id, hard_score, soft_score)

            # Clear any prior lesson rows (re-solves overwrite).
            await conn.execute("DELETE FROM tt_lesson WHERE solution_id = $1", solution_id)

            for lesson in timetable.lessons:
                if lesson.timeslot is None:
                    continue
                await conn.execute("""
                    INSERT INTO tt_lesson
                        (solution_id, class_id, subject_id, teacher_id, room_id,
                         day_of_week, period_index, pinned)
                    VALUES ($1, $2::uuid, $3::uuid, $4::uuid, $5::uuid, $6, $7, $8)
                """,
                    solution_id,
                    lesson.school_class.id,
                    lesson.subject.id,
                    lesson.teacher.id,
                    lesson.room.id if lesson.room else None,
                    lesson.timeslot.day_of_week,
                    lesson.timeslot.period_index,
                    lesson.pinned,
                )


async def mark_failed(pool: asyncpg.Pool, solution_id: str, error_message: str) -> None:
    async with pool.acquire() as conn:
        await conn.execute("""
            UPDATE tt_solution
            SET status = 'failed', error_message = $2, finished_at = NOW()
            WHERE id = $1
        """, solution_id, error_message)


async def mark_running(pool: asyncpg.Pool, solution_id: str) -> None:
    async with pool.acquire() as conn:
        await conn.execute("""
            UPDATE tt_solution SET status = 'running', started_at = NOW()
            WHERE id = $1
        """, solution_id)


async def mark_infeasible(pool: asyncpg.Pool, solution_id: str, hard_score: int, soft_score: int) -> None:
    async with pool.acquire() as conn:
        await conn.execute("""
            UPDATE tt_solution
            SET status = 'infeasible',
                hard_score = $2,
                soft_score = $3,
                finished_at = NOW()
            WHERE id = $1
        """, solution_id, hard_score, soft_score)