#!/usr/bin/env python3
"""
One-off script to generate example start_date/deadline data on shot tasks,
for testing/demoing the Project Schedule (Gantt) page.

Reuses the same shot-sampling approach as generate_note_example_data.py: the
dev database is heavily populated with disposable load-test fixtures, so this
picks a small, bounded sample of shots (from Dragon Quest) that have the
normal 4-task set, rather than scheduling all ~1400 shot tasks.

For each sampled shot, tasks are scheduled in pipeline order (layout ->
animation -> simulation -> lighting -> compositing, any other task types
appended after) with realistic overlapping durations, staggered around
today's date so the chart's "Today" marker falls inside the visible range.

Idempotent: skips any task that already has a start_date set, so it's safe
to re-run without clobbering manually-edited dates.
"""

from datetime import date, timedelta
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from database import DATABASE_URL, Base
from models.shot import Shot
from models.task import Task
from models.project import Project
import logging
import random

import models  # noqa: F401 - ensures every model is registered on Base.metadata

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

PROJECT_NAME = "Dragon Quest"
SAMPLE_SIZE = 15
RNG_SEED = 42

# Canonical pipeline order used to sequence a shot's tasks; anything not
# listed here keeps its natural (query) order, appended at the end.
PIPELINE_ORDER = ["layout", "animation", "simulation", "lighting", "compositing", "previz"]


def pipeline_index(task_type: str) -> int:
    return PIPELINE_ORDER.index(task_type) if task_type in PIPELINE_ORDER else len(PIPELINE_ORDER)


def generate_schedule_example_data():
    engine = create_engine(DATABASE_URL)
    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
    Base.metadata.create_all(bind=engine)

    db = SessionLocal()
    rng = random.Random(RNG_SEED)
    today = date.today()

    try:
        project = db.query(Project).filter(Project.name == PROJECT_NAME).first()
        if not project:
            logger.error(f"Project '{PROJECT_NAME}' not found")
            return
        logger.info(f"Project '{PROJECT_NAME}' (id={project.id})")

        # Shots with the normal 4-task set (excludes shots with zero tasks
        # and the handful of outliers with extra bulk-test tasks piled on).
        shots = db.query(Shot).filter(
            Shot.project_id == project.id, Shot.deleted_at.is_(None)
        ).order_by(Shot.id).all()
        candidates = [
            s for s in shots
            if db.query(Task).filter(Task.shot_id == s.id, Task.deleted_at.is_(None)).count() == 4
        ]
        sample_shots = candidates[:SAMPLE_SIZE]
        logger.info(f"{len(candidates)} candidate shots with a 4-task set, sampling {len(sample_shots)}")

        tasks_scheduled = 0
        tasks_skipped = 0

        for shot in sample_shots:
            tasks = db.query(Task).filter(
                Task.shot_id == shot.id, Task.deleted_at.is_(None)
            ).all()
            tasks.sort(key=lambda t: pipeline_index(t.task_type))

            # Stagger each shot's pipeline start across a window that straddles
            # today, so the demo chart shows a mix of past/current/future work.
            shot_start = today + timedelta(days=rng.randint(-20, 20))
            cursor = shot_start

            for task in tasks:
                if task.start_date is not None:
                    tasks_skipped += 1
                    # Still advance the cursor so later tasks in this shot
                    # don't bunch up if an earlier one was already scheduled.
                    cursor = task.deadline or cursor
                    continue

                duration = rng.randint(4, 12)
                task.start_date = cursor
                task.deadline = cursor + timedelta(days=duration)
                tasks_scheduled += 1

                # Next task starts a little before this one ends (realistic
                # pipeline overlap) rather than strictly back-to-back.
                cursor = task.deadline - timedelta(days=rng.randint(0, 4))

        db.commit()
        logger.info(f"Tasks scheduled: {tasks_scheduled}, skipped (already scheduled): {tasks_skipped}")

    except Exception as e:
        logger.error(f"Generation failed: {e}")
        db.rollback()
        raise
    finally:
        db.close()


if __name__ == "__main__":
    logger.info("Generating schedule example data...")
    generate_schedule_example_data()
    logger.info("Done!")
