#!/usr/bin/env python3
"""
One-off script to generate example Submission (and Review) data for
testing/demoing the work-submission and review flow.

Reuses the same shot-sampling approach as generate_note_example_data.py and
generate_schedule_example_data.py: picks the same bounded sample of Dragon
Quest shots (first 15 candidates with the normal 4-task set) rather than
touching all ~1400 shot tasks, so submissions line up with the tasks that
already have start_date/deadline from the schedule seed data.

Not every sampled task gets a submission (a "not started" task realistically
has none) - about 70% do, with 1-3 versions each. Earlier versions are given
a "retake" review, and the latest version is left pending, approved, or
retake'd, with the task's status updated to match. Each submission gets a
real placeholder file on disk under uploads/submissions/<task_id>/ (a tiny
real image for .jpg/.png so thumbnail generation works, plain bytes for
other formats) so the review UI has something to open instead of a 404.

Idempotent: skips any task that already has a submission, so it's safe to
re-run without creating duplicates.
"""

from datetime import datetime, timedelta, timezone
from pathlib import Path
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, Submission, Review, ReviewDecision
from models.project import Project, ProjectMember
from models.user import User, UserRole
from utils.file_handler import file_handler
from PIL import Image
import io
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  # same sample window as generate_schedule_example_data.py
RNG_SEED = 42
SUBMISSION_CHANCE = 0.7  # fraction of sampled tasks that get any submission

# Plausible deliverable formats per task type (all within FileHandler's
# SUPPORTED_FORMATS). Kept short - this is placeholder data, not a real
# pipeline convention.
TASK_TYPE_EXTENSIONS = {
    "layout": [".ma", ".mov"],
    "animation": [".mov", ".ma"],
    "simulation": [".mov", ".ma"],
    "lighting": [".exr", ".mov"],
    "compositing": [".mov", ".png", ".exr"],
    "previz": [".mov"],
}
DEFAULT_EXTENSIONS = [".mov"]

SUBMISSION_NOTES = [
    "First pass, ready for review.",
    "Addressed previous feedback, please take another look.",
    "Still WIP on some details but wanted to get eyes on it early.",
    "Final polish pass, should be good to approve.",
    None,  # some submissions have no note at all
]

RETAKE_FEEDBACK = [
    "Good start, but needs another pass before this is ready - see notes on the task.",
    "Close, but a few issues need fixing before this can move forward.",
    "Not quite there yet, please revise and resubmit.",
]

APPROVE_FEEDBACK = [
    "Looks great, approved!",
    "Nice work, this is ready to move on.",
    "Approved - matches the brief.",
]

IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png"}


def make_placeholder_file(task: Task, extension: str, version: int, rng: random.Random) -> tuple[str, int]:
    """Write a small real file to disk and return (relative_path, size)."""
    task_dir = file_handler.create_directory_structure(task.id, "submission")
    base_name = f"{(task.shot_id and 'shot' or 'asset')}_{task.task_type}_v{version:03d}{extension}"
    filename = file_handler.generate_unique_filename(base_name, f"v{version:03d}")
    file_path = task_dir / filename

    if extension in IMAGE_EXTENSIONS:
        # A tiny but real image so thumbnail generation (PIL.Image.open) works.
        color = (rng.randint(30, 220), rng.randint(30, 220), rng.randint(30, 220))
        img = Image.new("RGB", (64, 64), color)
        buf = io.BytesIO()
        img.save(buf, format="JPEG" if extension in (".jpg", ".jpeg") else "PNG")
        file_path.write_bytes(buf.getvalue())
    else:
        file_path.write_bytes(f"Placeholder {extension} submission file for testing.\n".encode())

    relative_path = file_handler.store_relative_path(str(file_path))
    return relative_path, file_path.stat().st_size


def generate_submission_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)
    now = datetime.now(timezone.utc)

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

        member_ids = [m.user_id for m in db.query(ProjectMember).filter(
            ProjectMember.project_id == project.id
        ).all()]
        if not member_ids:
            logger.error(f"No members found on project '{PROJECT_NAME}'")
            return

        reviewer_ids = [u.id for u in db.query(User).join(
            ProjectMember, ProjectMember.user_id == User.id
        ).filter(
            ProjectMember.project_id == project.id,
            User.role.in_([UserRole.DIRECTOR, UserRole.COORDINATOR])
        ).all()] or member_ids
        logger.info(f"Project '{PROJECT_NAME}' (id={project.id}), {len(member_ids)} members, {len(reviewer_ids)} potential reviewers")

        # 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_seen = 0
        tasks_skipped = 0
        tasks_submitted = 0
        submissions_created = 0
        reviews_created = 0

        for shot in sample_shots:
            tasks = db.query(Task).filter(
                Task.shot_id == shot.id, Task.deleted_at.is_(None)
            ).all()

            for task in tasks:
                tasks_seen += 1
                existing = db.query(Submission).filter(Submission.task_id == task.id).first()
                if existing:
                    tasks_skipped += 1
                    continue

                if rng.random() >= SUBMISSION_CHANCE:
                    continue

                extensions = TASK_TYPE_EXTENSIONS.get(task.task_type, DEFAULT_EXTENSIONS)
                version_count = rng.choice([1, 1, 2, 3])
                submitter_id = task.assigned_user_id or rng.choice(member_ids)

                # Spread submissions over the days leading up to the deadline
                # (or today, if the task has none) so submitted_at reads as a
                # real history rather than everything happening "now".
                anchor = (
                    datetime.combine(task.deadline, datetime.min.time(), tzinfo=timezone.utc)
                    if task.deadline else now
                )
                submitted_at = anchor - timedelta(days=version_count * 2)

                last_submission = None
                last_decision = None

                for version in range(1, version_count + 1):
                    extension = rng.choice(extensions)
                    relative_path, file_size = make_placeholder_file(task, extension, version, rng)
                    thumbnail_path = None
                    if extension in IMAGE_EXTENSIONS:
                        thumbnail_path = file_handler.create_thumbnail(relative_path)

                    submission = Submission(
                        task_id=task.id,
                        user_id=submitter_id,
                        file_path=relative_path,
                        file_name=Path(relative_path).name,
                        version_number=version,
                        notes=rng.choice(SUBMISSION_NOTES),
                        submitted_at=submitted_at
                    )
                    db.add(submission)
                    db.flush()  # need submission.id for a Review FK
                    submissions_created += 1
                    submitted_at += timedelta(days=2)

                    is_last_version = version == version_count
                    if not is_last_version:
                        # Earlier versions were superseded because they needed work.
                        review = Review(
                            submission_id=submission.id,
                            reviewer_id=rng.choice(reviewer_ids),
                            decision=ReviewDecision.RETAKE,
                            feedback=rng.choice(RETAKE_FEEDBACK),
                            reviewed_at=submission.submitted_at + timedelta(hours=rng.randint(2, 20))
                        )
                        db.add(review)
                        reviews_created += 1
                    else:
                        # Latest version: leave some pending review, resolve the rest.
                        outcome = rng.choice(["pending", "approved", "approved", "retake"])
                        if outcome != "pending":
                            decision = ReviewDecision.APPROVED if outcome == "approved" else ReviewDecision.RETAKE
                            feedback = rng.choice(APPROVE_FEEDBACK if outcome == "approved" else RETAKE_FEEDBACK)
                            review = Review(
                                submission_id=submission.id,
                                reviewer_id=rng.choice(reviewer_ids),
                                decision=decision,
                                feedback=feedback,
                                reviewed_at=submission.submitted_at + timedelta(hours=rng.randint(2, 20))
                            )
                            db.add(review)
                            reviews_created += 1
                        last_decision = outcome

                    last_submission = submission

                if last_submission:
                    task.status = {
                        "approved": "approved",
                        "retake": "retake",
                        "pending": "submitted"
                    }[last_decision]
                    tasks_submitted += 1

        db.commit()
        logger.info(f"Tasks seen: {tasks_seen}, skipped (already had a submission): {tasks_skipped}")
        logger.info(f"Tasks given submissions: {tasks_submitted}")
        logger.info(f"Submissions created: {submissions_created}, reviews created: {reviews_created}")

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


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