Aislo/common_util/common_util_workflow_state.py

"""워크플로우 단계별 상태 관리를 위한 공통 유틸리티."""

import json
from datetime import datetime
from typing import Any, Dict

import aiomysql

STAGE_KEYS = [
    "FILE_INPUT",  # 0
    "WF1_SURFACE",  # 1
    "WF2_ROUTE",  # 2
    "WF3_PROFILE_CROSS",  # 3
    "WF4_DESIGN_DETAIL",  # 4
    "WF5_QUANTITY",  # 5
    "WF6_ESTIMATION",  # 6
]


async def initialize_project_stages(cursor: aiomysql.DictCursor, project_id: str) -> None:
    """프로젝트 생성 시 7개의 단계를 NOT_STARTED 상태로 시드한다."""
    for stage_no, stage_key in enumerate(STAGE_KEYS):
        await cursor.execute(
            """
            INSERT INTO project_workflow_stages (
                project_id, stage_no, stage_key, state, progress_percent
            )
            VALUES (%s, %s, %s, 'NOT_STARTED', 0)
            ON DUPLICATE KEY UPDATE
                state = 'NOT_STARTED',
                progress_percent = 0,
                params = NULL,
                message = NULL,
                started_at = NULL,
                completed_at = NULL
            """,
            (project_id, stage_no, stage_key),
        )


async def start_stage(
    cursor: aiomysql.DictCursor,
    project_id: str,
    stage_no: int,
    params: Dict[str, Any] | None = None,
) -> None:
    """단계를 시작하여 IN_PROGRESS 상태로 만들고, 이후 단계들을 STALE로 전환한다."""
    now = datetime.utcnow()
    params_json = json.dumps(params, ensure_ascii=False) if params is not None else None

    # 해당 단계 시작
    await cursor.execute(
        """
        UPDATE project_workflow_stages
        SET state = 'IN_PROGRESS',
            progress_percent = 0,
            params = COALESCE(%s, params),
            message = NULL,
            started_at = %s,
            completed_at = NULL
        WHERE project_id = %s AND stage_no = %s
        """,
        (params_json, now, project_id, stage_no),
    )

    # 역방향 재작업 무효화 (stale 전파): 이후 단계 중 COMPLETE인 것들을 STALE로 변경
    await cursor.execute(
        """
        UPDATE project_workflow_stages
        SET state = 'STALE'
        WHERE project_id = %s AND stage_no > %s AND state = 'COMPLETE'
        """,
        (project_id, stage_no),
    )

    # projects.status 캐시 업데이트 (호환성 유지)
    status_str = f"WF{stage_no}_ANALYZING" if stage_no > 0 else "FILE_INPUT_PROCESSING"
    await cursor.execute(
        """
        UPDATE projects
        SET status = %s, updated_at = %s
        WHERE id = %s
        """,
        (status_str, now, project_id),
    )


async def complete_stage(cursor: aiomysql.DictCursor, project_id: str, stage_no: int) -> None:
    """단계를 완료 상태로 전환한다."""
    now = datetime.utcnow()
    await cursor.execute(
        """
        UPDATE project_workflow_stages
        SET state = 'COMPLETE',
            progress_percent = 100,
            completed_at = %s
        WHERE project_id = %s AND stage_no = %s
        """,
        (now, project_id, stage_no),
    )

    # projects.status 캐시 업데이트 (호환성 유지)
    status_str = f"WF{stage_no}_COMPLETE" if stage_no > 0 else "FILE_UPLOADED"
    if stage_no == 6:
        status_str = "DONE"

    await cursor.execute(
        """
        UPDATE projects
        SET status = %s, updated_at = %s
        WHERE id = %s
        """,
        (status_str, now, project_id),
    )


async def update_stage_progress(
    cursor: aiomysql.DictCursor, project_id: str, stage_no: int, progress: int
) -> None:
    """단계 진행률을 업데이트한다."""
    await cursor.execute(
        """
        UPDATE project_workflow_stages
        SET progress_percent = %s
        WHERE project_id = %s AND stage_no = %s
        """,
        (progress, project_id, stage_no),
    )


async def fail_stage(
    cursor: aiomysql.DictCursor, project_id: str, stage_no: int, message: str
) -> None:
    """단계를 실패 상태로 전환한다."""
    now = datetime.utcnow()
    await cursor.execute(
        """
        UPDATE project_workflow_stages
        SET state = 'FAILED',
            message = %s
        WHERE project_id = %s AND stage_no = %s
        """,
        (message, project_id, stage_no),
    )

    # projects.status 캐시 업데이트
    status_str = f"WF{stage_no}_FAILED" if stage_no > 0 else "FILE_INPUT_FAILED"
    await cursor.execute(
        """
        UPDATE projects
        SET status = %s, updated_at = %s
        WHERE id = %s
        """,
        (status_str, now, project_id),
    )


async def get_workflow_state(cursor: aiomysql.DictCursor, project_id: str) -> Dict[str, Any]:
    """프로젝트의 모든 단계 상태를 조회하여 요약 및 배열로 반환한다."""
    await cursor.execute(
        """
        SELECT stage_no, stage_key, state, progress_percent, params, message,
               DATE_FORMAT(started_at, '%%Y-%%m-%%dT%%H:%%i:%%s') AS started_at,
               DATE_FORMAT(completed_at, '%%Y-%%m-%%dT%%H:%%i:%%s') AS completed_at
        FROM project_workflow_stages
        WHERE project_id = %s
        ORDER BY stage_no ASC
        """,
        (project_id,),
    )
    rows = await cursor.fetchall()

    if not rows:
        return {"project_id": project_id, "current_stage": 0, "stages": []}

    stages_list = []
    for r in rows:
        params_val = None
        if r.get("params"):
            try:
                params_val = (
                    json.loads(r["params"]) if isinstance(r["params"], str) else r["params"]
                )
            except Exception:
                params_val = r["params"]

        stages_list.append(
            {
                "stage_no": r["stage_no"],
                "stage_key": r["stage_key"],
                "state": r["state"],
                "progress_percent": r["progress_percent"],
                "params": params_val,
                "message": r["message"],
                "started_at": r["started_at"],
                "completed_at": r["completed_at"],
            }
        )

    current_stage = 0
    for stage in stages_list:
        if stage["stage_no"] == 0:
            continue
        prev_stage = stages_list[stage["stage_no"] - 1]
        if prev_stage["state"] == "COMPLETE":
            current_stage = stage["stage_no"]
        else:
            break

    return {"project_id": project_id, "current_stage": current_stage, "stages": stages_list}