video-gen/video-gen-api/app/services/module_generation_step_common_service.py

from __future__ import annotations

import base64
import json
from copy import deepcopy
from datetime import datetime, timezone
from typing import Any

from sqlalchemy.orm.attributes import flag_modified

from app.config import settings
from app.models.chat_generation_task import ChatGenerationTask
from app.services.resource_signed_url_service import build_resource_signed_url

STEP_IO_SCHEMA_VERSION = "module_generation_step_io_v1"


def utc_now() -> datetime:
    return datetime.now(timezone.utc)


def json_dumps(data: Any) -> str | None:
    if data is None:
        return None
    return json.dumps(data, ensure_ascii=False, default=str)


def parse_json(value: Any, fallback: Any = None) -> Any:
    if value is None or value == "":
        return fallback
    if isinstance(value, (dict, list)):
        return value
    if isinstance(value, str):
        try:
            return json.loads(value)
        except Exception:
            return fallback
    return fallback


def build_step_input(
    *,
    step_code: str,
    payload: dict[str, Any] | None = None,
    source_step_id: str | None = None,
    parent_step_id: str | None = None,
    context: dict[str, Any] | None = None,
    schema_version: str = STEP_IO_SCHEMA_VERSION,
) -> dict[str, Any]:
    return {
        "schema_version": schema_version,
        "step_code": step_code,
        "source": {
            "source_step_id": source_step_id,
            "parent_step_id": parent_step_id,
        },
        "payload": payload or {},
        "context": context or {},
    }


def build_step_output(
    *,
    step_code: str,
    status: str,
    payload: dict[str, Any] | None = None,
    result: dict[str, Any] | None = None,
    usage: dict[str, Any] | None = None,
    error: dict[str, Any] | None = None,
    schema_version: str = STEP_IO_SCHEMA_VERSION,
) -> dict[str, Any]:
    return {
        "schema_version": schema_version,
        "step_code": step_code,
        "status": status,
        "payload": payload or {},
        "result": result or {},
        "usage": usage or {},
        "error": error or {},
    }


def is_wrapped_step_io(value: Any, *, schema_version: str = STEP_IO_SCHEMA_VERSION) -> bool:
    return isinstance(value, dict) and value.get("schema_version") == schema_version


def step_payload(value: Any, *, schema_version: str = STEP_IO_SCHEMA_VERSION) -> dict[str, Any]:
    data = parse_json(value, {}) or {}
    if is_wrapped_step_io(data, schema_version=schema_version):
        payload = data.get("payload")
        return payload if isinstance(payload, dict) else {}
    return data if isinstance(data, dict) else {}


def step_result(value: Any, *, schema_version: str = STEP_IO_SCHEMA_VERSION) -> dict[str, Any]:
    data = parse_json(value, {}) or {}
    if is_wrapped_step_io(data, schema_version=schema_version):
        result = data.get("result")
        if isinstance(result, dict) and result:
            return result
        payload = data.get("payload")
        return payload if isinstance(payload, dict) else {}
    return data if isinstance(data, dict) else {}


def step_usage(value: Any, *, schema_version: str = STEP_IO_SCHEMA_VERSION) -> dict[str, Any]:
    data = parse_json(value, {}) or {}
    if is_wrapped_step_io(data, schema_version=schema_version):
        usage = data.get("usage")
        return usage if isinstance(usage, dict) else {}
    usage = data.get("token_usage") if isinstance(data, dict) else {}
    return usage if isinstance(usage, dict) else {}


def unwrap_step_output(value: Any, *, schema_version: str = STEP_IO_SCHEMA_VERSION) -> dict[str, Any]:
    data = parse_json(value, {}) or {}
    if not is_wrapped_step_io(data, schema_version=schema_version):
        return data if isinstance(data, dict) else {}
    merged: dict[str, Any] = {}
    payload = data.get("payload")
    result = data.get("result")
    usage = data.get("usage")
    if isinstance(payload, dict):
        merged.update(payload)
    if isinstance(result, dict):
        merged.update(result)
    if isinstance(usage, dict) and usage:
        merged["token_usage"] = usage
    return merged


def merge_dict(old: dict[str, Any] | None, new: dict[str, Any] | None) -> dict[str, Any]:
    merged = dict(old or {})
    for key, value in (new or {}).items():
        if value is not None:
            merged[key] = value
    return merged


def force_set_json(model_obj: Any, field_name: str, value: Any) -> None:
    """强制持久化 JSON / JSONB 字段。

    SQLAlchemy 对 dict/list 的嵌套原地修改不会稳定触发 dirty 判定。
    所有编辑类接口在写入 input_json / output_json 时统一走这里：
    1. deepcopy 断开旧引用；
    2. 整体重新赋值；
    3. flag_modified 显式标记字段已变更。
    """
    setattr(model_obj, field_name, deepcopy(value))
    flag_modified(model_obj, field_name)


def snapshot_from_chat(chat_task: ChatGenerationTask | None) -> dict[str, Any]:
    if not chat_task:
        return {}
    return {"chat_task_id": chat_task.id, "status": chat_task.status, "pipeline_stage": chat_task.pipeline_stage}


def build_file_url_or_data_uri(file_url: str) -> str:
    if file_url.startswith("http://") or file_url.startswith("https://") or file_url.startswith("data:"):
        return file_url
    file_url_sign = build_resource_signed_url(resource_url=file_url, expire_seconds=86400)
    return f"{settings.BASE_URL}{file_url_sign}"
    # try:
    #     with open(file_url, "rb") as f:
    #         return "data:image/png;base64," + base64.b64encode(f.read()).decode("utf-8")
    # except Exception:
    #     return f"{settings.STORAGE_BASE_URL.rstrip('/')}/{file_url.lstrip('/')}"