"""현재 state부터 다음 게이트까지 단계를 이어 돌림
어느 잡을 언제 돌릴지는 worker가 정하고 여기는 실행만 함
"""
import re
import time
import traceback

from sqlalchemy.ext.asyncio import AsyncSession

from models.pipeline_state import PlayreelState, PosterAliveState, stage_keys
from pipelines import playreel, poster_alive
from pipelines.gate import PLAYREEL_GATE_AFTER, POSTER_ALIVE_GATE_AFTER
from tables.task import PlayreelTask, PosterAliveTask

QUEUED = "queued"
RUNNING = "running"
AWAITING_REVIEW = "awaiting_review"
FAILED = "failed"
DONE = "done"

POSTER_ALIVE_RUNNERS = {
    PosterAliveState.DETECT: poster_alive.run_detect,
    PosterAliveState.NARRATION_TEXT: poster_alive.run_narration_text,
    PosterAliveState.MOTION: poster_alive.run_motion,
    PosterAliveState.TTS: poster_alive.run_tts,
    PosterAliveState.BGM: poster_alive.run_bgm,
    PosterAliveState.I2V: poster_alive.run_i2v,
    PosterAliveState.RENDER: poster_alive.run_render,
}

PLAYREEL_RUNNERS = {
    PlayreelState.FETCH: playreel.run_fetch,
    PlayreelState.SPLIT: playreel.run_split,
    PlayreelState.UPSCALE: playreel.run_upscale,
    PlayreelState.ANALYZE: playreel.run_analyze,
    PlayreelState.MOTION: playreel.run_motion,
    PlayreelState.NARRATION: playreel.run_narration,
    PlayreelState.TTS: playreel.run_tts,
    PlayreelState.BGM: playreel.run_bgm,
    PlayreelState.I2V: playreel.run_i2v,
    PlayreelState.HYBRID: playreel.run_hybrid,
    PlayreelState.COMPOSE: playreel.run_compose,
    # ⑫ review는 대응하는 서비스가 없어 통과만 시킨다
}

STATES = {PosterAliveTask: PosterAliveState, PlayreelTask: PlayreelState}
RUNNERS = {PosterAliveTask: POSTER_ALIVE_RUNNERS, PlayreelTask: PLAYREEL_RUNNERS}


def mark(task, stage: str, status: str) -> None:
    timings = dict(task.stage_timings or {})
    entry = dict(timings.get(stage) or {"started": None, "ended": None})
    entry["status"] = status
    if status == RUNNING:
        entry["started"] = time.time()
    elif status in ("done", FAILED):
        entry["ended"] = time.time()
    timings[stage] = entry
    task.stage_timings = timings


def reset_from(task, stage: str) -> None:
    """되돌리기·재시도에서 그 단계부터 다시 돌게 만든다"""
    state_type = STATES[type(task)]
    keys = stage_keys(state_type)
    timings = dict(task.stage_timings or {})
    for key in keys[keys.index(stage):]:
        timings[key] = {"status": "idle", "started": None, "ended": None}
    task.stage_timings = timings
    task.state = state_type(stage)


URL_IN_TEXT = re.compile(r"https?://\S+")
MAX_ERROR_DETAIL = 500


def safe_detail(failure: BaseException) -> str:
    """검수 화면까지 가는 값이라 트레이스백과 스토리지 주소는 뺀다"""
    message = f"{type(failure).__name__}: {failure}"
    return URL_IN_TEXT.sub("<url>", message)[:MAX_ERROR_DETAIL]


def stops_here(task, stage) -> str | None:
    """이 단계 뒤에 사람이 볼 게이트가 서는가. playreel은 게이트 이름을 돌려준다"""
    if isinstance(task, PlayreelTask):
        gate = PLAYREEL_GATE_AFTER.get(stage)
        return gate.value if gate else None
    if stage == POSTER_ALIVE_GATE_AFTER and not task.skip_review:
        return AWAITING_REVIEW
    return None


async def step(session: AsyncSession, task) -> bool:
    """한 단계를 돌린다. 게이트에 닿으면 False를 돌려 멈춤을 알린다"""
    state_type = STATES[type(task)]
    stage = state_type(task.state)
    runner = RUNNERS[type(task)].get(stage)

    if runner is None:                       # 러너 없는 마지막 단계
        mark(task, stage.value, "done")
        task.state = state_type.COMPLETED
        await session.commit()
        return True

    mark(task, stage.value, RUNNING)
    await session.commit()

    await runner(session, task)              # state를 다음 단계로 올리고 커밋한다

    mark(task, stage.value, "done")
    gate = stops_here(task, stage)
    if gate:
        task.status = AWAITING_REVIEW
        if isinstance(task, PlayreelTask):
            task.gate = gate
        await session.commit()
        return False
    await session.commit()
    return True


async def run_until_gate(session: AsyncSession, task) -> None:
    """다음 게이트나 끝까지 민다. 실패는 잡에 기록하고 조용히 돌아온다"""
    state_type = STATES[type(task)]
    try:
        while task.state != state_type.COMPLETED:
            if not await step(session, task):
                return
    except Exception as failure:
        failed_stage = str(task.state)
        traceback.print_exc()                # 전문은 서버 로그에만 남긴다
        await session.rollback()
        await session.refresh(task)
        mark(task, failed_stage, FAILED)
        task.status = FAILED
        task.error_stage = failed_stage
        task.error_detail = safe_detail(failure)
        await session.commit()
        return

    task.status = DONE
    if isinstance(task, PlayreelTask):
        task.gate = None
        task.version += 1
    await session.commit()