"""add story generation status fields

Revision ID: 0009_add_story_generation_statuses
Revises: 0008_add_pages_to_stories
Create Date: 2026-04-17

"""

import sqlalchemy as sa
from alembic import op


# revision identifiers, used by Alembic.
revision = "0009_add_story_generation_statuses"
down_revision = "0008_add_pages_to_stories"
branch_labels = None
depends_on = None


stories = sa.table(
    "stories",
    sa.column("id", sa.Integer),
    sa.column("story_text", sa.Text),
    sa.column("pages", sa.JSON),
    sa.column("cover_prompt", sa.Text),
    sa.column("image_url", sa.String(length=500)),
    sa.column("generation_status", sa.String(length=32)),
    sa.column("image_status", sa.String(length=32)),
    sa.column("audio_status", sa.String(length=32)),
)


def _resolve_image_status(row: dict) -> str:
    pages = row.get("pages") or []

    expected_assets = 0
    ready_assets = 0

    if row.get("cover_prompt") or row.get("image_url"):
        expected_assets += 1
        if row.get("image_url"):
            ready_assets += 1

    for page in pages:
        if not isinstance(page, dict):
            continue
        if not page.get("image_prompt") and not page.get("image_url"):
            continue
        expected_assets += 1
        if page.get("image_url"):
            ready_assets += 1

    if expected_assets == 0:
        return "not_requested"

    if ready_assets == expected_assets:
        return "ready"

    return "failed"


def _resolve_generation_status(
    *,
    story_text: str | None,
    pages: list[dict] | None,
    image_status: str,
    audio_status: str,
) -> str:
    has_narrative = bool(story_text) or bool(pages)
    if not has_narrative:
        return "failed"

    if "generating" in {image_status, audio_status}:
        return "assets_generating"

    if "failed" in {image_status, audio_status}:
        return "degraded_completed"

    if image_status == "not_requested" and audio_status == "not_requested":
        return "narrative_ready"

    return "completed"


def upgrade() -> None:
    op.add_column(
        "stories",
        sa.Column(
            "generation_status",
            sa.String(length=32),
            nullable=False,
            server_default="narrative_ready",
        ),
    )
    op.add_column(
        "stories",
        sa.Column(
            "image_status",
            sa.String(length=32),
            nullable=False,
            server_default="not_requested",
        ),
    )
    op.add_column(
        "stories",
        sa.Column(
            "audio_status",
            sa.String(length=32),
            nullable=False,
            server_default="not_requested",
        ),
    )
    op.add_column("stories", sa.Column("last_error", sa.Text(), nullable=True))

    connection = op.get_bind()
    rows = connection.execute(
        sa.select(
            stories.c.id,
            stories.c.story_text,
            stories.c.pages,
            stories.c.cover_prompt,
            stories.c.image_url,
        )
    ).mappings()

    for row in rows:
        image_status = _resolve_image_status(row)
        audio_status = "not_requested"
        generation_status = _resolve_generation_status(
            story_text=row.get("story_text"),
            pages=row.get("pages"),
            image_status=image_status,
            audio_status=audio_status,
        )

        connection.execute(
            stories.update()
            .where(stories.c.id == row["id"])
            .values(
                generation_status=generation_status,
                image_status=image_status,
                audio_status=audio_status,
            )
        )


def downgrade() -> None:
    op.drop_column("stories", "last_error")
    op.drop_column("stories", "audio_status")
    op.drop_column("stories", "image_status")
    op.drop_column("stories", "generation_status")