"""Lightweight generation job/event tracking."""

from __future__ import annotations

from typing import Any

from sqlalchemy.ext.asyncio import AsyncSession

from app.db.models import GenerationJob, GenerationJobEvent, Story


def _story_snapshot(story: Story | None) -> dict[str, Any]:
    if story is None:
        return {}

    return {
        "story_id": story.id,
        "mode": story.mode,
        "generation_status": story.generation_status,
        "image_status": story.image_status,
        "audio_status": story.audio_status,
        "retryable_assets": story.retryable_assets,
        "last_error": story.last_error,
    }


def _job_status_from_story(story: Story) -> str:
    if story.generation_status == "failed":
        return "failed"
    if story.generation_status == "degraded_completed":
        return "degraded_completed"
    return "completed"


async def create_generation_job(
    db: AsyncSession,
    *,
    user_id: str,
    output_mode: str,
    input_type: str,
    request_payload: dict[str, Any],
    story_id: int | None = None,
) -> GenerationJob:
    """Create a generation job and record its first event."""

    job = GenerationJob(
        user_id=user_id,
        story_id=story_id,
        output_mode=output_mode,
        input_type=input_type,
        status="running",
        current_step="request_accepted",
        request_payload=request_payload,
        result_snapshot={},
    )
    db.add(job)
    await db.flush()
    await record_generation_event(
        db,
        job=job,
        story_id=story_id,
        event_type="request_accepted",
        status="succeeded",
        message="Generation request accepted.",
        metadata={"output_mode": output_mode, "input_type": input_type},
        commit=False,
    )
    await db.commit()
    await db.refresh(job)
    return job


async def record_generation_event(
    db: AsyncSession,
    *,
    job: GenerationJob,
    event_type: str,
    status: str,
    story_id: int | None = None,
    message: str | None = None,
    metadata: dict[str, Any] | None = None,
    commit: bool = True,
) -> GenerationJobEvent:
    """Append one event to an existing generation job."""

    event = GenerationJobEvent(
        job_id=job.id,
        story_id=story_id if story_id is not None else job.story_id,
        event_type=event_type,
        status=status,
        message=message,
        event_metadata=metadata or {},
    )
    db.add(event)
    if commit:
        await db.commit()
    return event


async def finish_generation_job(
    db: AsyncSession,
    *,
    job: GenerationJob,
    story: Story | None,
    status: str | None = None,
    current_step: str,
    error_message: str | None = None,
    message: str | None = None,
    metadata: dict[str, Any] | None = None,
) -> GenerationJob:
    """Mark a generation job as completed/degraded/failed and append a final event."""

    job.story_id = story.id if story is not None else job.story_id
    job.status = status or (_job_status_from_story(story) if story is not None else "failed")
    job.current_step = current_step
    job.error_message = error_message
    job.result_snapshot = _story_snapshot(story)
    await record_generation_event(
        db,
        job=job,
        story_id=job.story_id,
        event_type=current_step,
        status=job.status,
        message=message,
        metadata={
            **(metadata or {}),
            "result_snapshot": job.result_snapshot,
        },
        commit=False,
    )
    await db.commit()
    await db.refresh(job)
    return job