PYTV/core/services/scheduler.py

"""
Schedule generator — respects ChannelSourceRule assignments.

Source selection priority:
  1. If any rules with rule_mode='prefer' exist, items from those sources
     are weighted much more heavily.
  2. Items from rule_mode='allow' sources fill the rest.
  3. Items from rule_mode='avoid' sources are only used as a last resort
     (weight × 0.1).
  4. Items from rule_mode='block' sources are NEVER scheduled.
  5. If NO ChannelSourceRule rows exist for this channel, falls back to
     the old behaviour (all items in the channel's library).
"""

import random
import uuid
from datetime import datetime, timedelta, date, timezone

from core.models import (
    Channel, ChannelSourceRule, ScheduleTemplate,
    ScheduleBlock, Airing, MediaItem,
)


class ScheduleGenerator:
    """
    Reads ScheduleTemplate + ScheduleBlocks for a channel and fills the day
    with concrete Airing rows, picking MediaItems according to the channel's
    ChannelSourceRule assignments.
    """

    def __init__(self, channel: Channel):
        self.channel = channel

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def generate_for_date(self, target_date: date) -> int:
        """
        Idempotent generation of airings for `target_date`.
        Returns the number of new Airing rows created.
        """
        template = self._get_template()
        if not template:
            return 0

        target_weekday_bit = 1 << target_date.weekday()
        blocks = template.scheduleblock_set.all()
        airings_created = 0

        for block in blocks:
            if not (block.day_of_week_mask & target_weekday_bit):
                continue

            start_dt = datetime.combine(target_date, block.start_local_time, tzinfo=timezone.utc)
            end_dt   = datetime.combine(target_date, block.end_local_time,   tzinfo=timezone.utc)

            # Midnight-wrap support (e.g. 23:00–02:00)
            if end_dt <= start_dt:
                end_dt += timedelta(days=1)

            # Clear existing airings in this window (idempotency)
            Airing.objects.filter(
                channel=self.channel,
                starts_at__gte=start_dt,
                starts_at__lt=end_dt,
            ).delete()

            available_items = self._get_weighted_items(block)
            if not available_items:
                continue

            airings_created += self._fill_block(
                template, block, start_dt, end_dt, available_items
            )

        return airings_created

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _get_template(self):
        """Pick the highest-priority active ScheduleTemplate for this channel."""
        qs = ScheduleTemplate.objects.filter(
            channel=self.channel, is_active=True
        ).order_by('-priority')
        return qs.first()

    def _get_weighted_items(self, block: ScheduleBlock) -> list:
        """
        Build a weighted pool of MediaItems respecting ChannelSourceRule.

        Returns a flat list with items duplicated according to their effective
        weight (rounded to nearest int, min 1) so random.choice() gives the
        right probability distribution without needing numpy.
        """
        rules = list(
            ChannelSourceRule.objects.filter(channel=self.channel)
            .select_related('media_source')
        )

        if rules:
            # ── Rules exist: build filtered + weighted pool ───────────────
            allowed_source_ids  = set()  # allow + prefer
            blocked_source_ids  = set()  # block
            avoid_source_ids    = set()  # avoid
            source_weights: dict[int, float] = {}

            for rule in rules:
                sid = rule.media_source_id
                mode = rule.rule_mode
                w = float(rule.weight or 1.0)

                if mode == 'block':
                    blocked_source_ids.add(sid)
                elif mode == 'avoid':
                    avoid_source_ids.add(sid)
                    source_weights[sid] = w * 0.1   # heavily discounted
                elif mode == 'prefer':
                    allowed_source_ids.add(sid)
                    source_weights[sid] = w * 3.0   # boosted
                else:  # 'allow'
                    allowed_source_ids.add(sid)
                    source_weights[sid] = w

            # Build base queryset from allowed + avoid sources (not blocked)
            eligible_source_ids = (allowed_source_ids | avoid_source_ids) - blocked_source_ids

            if not eligible_source_ids:
                return []

            base_qs = MediaItem.objects.filter(
                media_source_id__in=eligible_source_ids,
                is_active=True,
            ).exclude(item_kind='bumper').select_related('media_source')

        else:
            # ── No rules: fall back to full library (old behaviour) ────────
            base_qs = MediaItem.objects.filter(
                media_source__library=self.channel.library,
                is_active=True,
            ).exclude(item_kind='bumper')
            source_weights = {}

        # Optionally filter by genre if block specifies one
        if block.default_genre:
            base_qs = base_qs.filter(genres=block.default_genre)

        items = list(base_qs)
        if not items:
            return []

        if not source_weights:
            # No weight information — plain shuffle
            random.shuffle(items)
            return items

        # Build weighted list: each item appears ⌈weight⌉ times
        weighted: list[MediaItem] = []
        for item in items:
            w = source_weights.get(item.media_source_id, 1.0)
            copies = max(1, round(w))
            weighted.extend([item] * copies)

        random.shuffle(weighted)
        return weighted

    def _fill_block(
        self,
        template: ScheduleTemplate,
        block: ScheduleBlock,
        start_dt: datetime,
        end_dt: datetime,
        items: list,
    ) -> int:
        """Fill start_dt→end_dt with sequential Airings, cycling through items."""
        cursor = start_dt
        idx = 0
        created = 0
        batch = uuid.uuid4()

        while cursor < end_dt:
            item = items[idx % len(items)]
            idx += 1

            duration = timedelta(seconds=max(item.runtime_seconds or 1800, 1))

            # Don't let a single item overshoot the end by more than its own length
            if cursor + duration > end_dt + timedelta(hours=1):
                break

            Airing.objects.create(
                channel=self.channel,
                schedule_template=template,
                schedule_block=block,
                media_item=item,
                starts_at=cursor,
                ends_at=cursor + duration,
                slot_kind="program",
                status="scheduled",
                source_reason="template",
                generation_batch_uuid=batch,
            )

            cursor += duration
            created += 1

        return created