PMFlow/backend/app/utils/seed_projects.py

"""Seed projects and signals, including Pixelle SSTC bootstrap."""

from __future__ import annotations

import asyncio
from typing import Optional

from sqlalchemy import func, or_, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import sessionmaker

from app.models import Project, ProjectSignalLink, Signal, TruthNode
from app.utils.seed_projects_helpers import (
    build_ai_signal,
    build_pixelle_signals,
    create_pixelle_charter_and_tasks,
    ensure_seed_user,
    pick_signal_ids,
)

PIXELLE_NAME = "Pixelle-Video 集成"


async def seed_pixelle_project(db: Optional[AsyncSession] = None):
    """创建或恢复 Pixelle-Video 示例项目与 SSTC 数据。"""
    from app.core.database import engine

    if db is None:
        async_session_maker = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
        async with async_session_maker() as session:
            await _seed_projects(session)
    else:
        await _seed_projects(db)


async def _ensure_pixelle_sstc_data(session: AsyncSession, pixelle_project: Project) -> bool:
    """为已有项目补齐关联信号、宪章和任务。"""
    pid = pixelle_project.id

    link_count = await session.execute(
        select(func.count()).select_from(ProjectSignalLink).where(ProjectSignalLink.project_id == pid)
    )
    has_links = (link_count.scalar() or 0) > 0

    charter_count = await session.execute(
        select(func.count()).select_from(TruthNode).where(
            TruthNode.project_id == pid,
            TruthNode.node_type == "charter",
        )
    )
    has_charters = (charter_count.scalar() or 0) > 0
    if has_links and has_charters:
        return False

    changed = False
    pixelle_signals: list[Signal] = []

    if not has_links:
        pixelle_signals = build_pixelle_signals()
        for signal in pixelle_signals:
            session.add(signal)
        await session.flush()

        for signal in pixelle_signals:
            session.add(
                ProjectSignalLink(
                    project_id=pid,
                    signal_id=signal.id,
                    link_type="related",
                    relevance_score=1.0,
                )
            )

        pixelle_project.signal_count = len(pixelle_signals)
        changed = True
        await session.flush()

    if not has_charters:
        seed_user = await ensure_seed_user(session)
        if pixelle_signals:
            signal_ids = pick_signal_ids(pixelle_signals)
        else:
            link_rows = await session.execute(
                select(ProjectSignalLink.signal_id).where(ProjectSignalLink.project_id == pid)
            )
            signal_ids = [row[0] for row in link_rows.fetchall()]

        await create_pixelle_charter_and_tasks(
            session,
            project_id=pid,
            created_by=seed_user.id,
            signal_ids=signal_ids,
        )

        changed = True
        print(f"[SSTC] 已为「{PIXELLE_NAME}」恢复宪章 1 个、任务 3 个")

    return changed


async def _seed_projects(session: AsyncSession):
    result = await session.execute(select(Project).where(or_(Project.title == PIXELLE_NAME, Project.name == PIXELLE_NAME)))
    existing = result.scalar_one_or_none()
    if existing:
        restored = await _ensure_pixelle_sstc_data(session, existing)
        if restored:
            await session.commit()
        else:
            print("示例项目已存在且 SSTC 数据完整，跳过")
        return

    pixelle_signals = build_pixelle_signals()
    for signal in pixelle_signals:
        session.add(signal)
    await session.flush()

    pixelle_project = Project(
        name=PIXELLE_NAME,
        title=PIXELLE_NAME,
        description="集成 Pixelle-Video AI 视频生成平台到主系统",
        status="active",
        assigned_role="backend",
        signal_count=len(pixelle_signals),
    )
    session.add(pixelle_project)
    await session.flush()

    for signal in pixelle_signals:
        session.add(
            ProjectSignalLink(
                project_id=pixelle_project.id,
                signal_id=signal.id,
                link_type="related",
                relevance_score=1.0,
            )
        )

    ai_analyst_project = Project(
        name="AI 分析引擎优化",
        title="AI 分析引擎优化",
        description="优化 LLM 服务性能，实现多模型投票机制和 RAG 检索增强",
        status="active",
        assigned_role="ai_engineer",
        signal_count=1,
    )
    session.add(ai_analyst_project)
    await session.flush()

    ai_signal = build_ai_signal()
    session.add(ai_signal)
    await session.flush()

    session.add(
        ProjectSignalLink(
            project_id=ai_analyst_project.id,
            signal_id=ai_signal.id,
            link_type="related",
            relevance_score=0.9,
        )
    )

    seed_user = await ensure_seed_user(session)
    await create_pixelle_charter_and_tasks(
        session,
        project_id=pixelle_project.id,
        created_by=seed_user.id,
        signal_ids=pick_signal_ids(pixelle_signals),
    )

    await session.commit()

    print("创建示例项目:")
    print(f"  - {pixelle_project.title} ({pixelle_project.assigned_role})")
    print(f"  - {ai_analyst_project.title} ({ai_analyst_project.assigned_role})")
    print(f"\n关联信号: {len(pixelle_signals) + 1} 条")
    print("[SSTC] Pixelle-Video 宪章 1 个，任务 3 个")


async def main():
    print("开始初始化示例数据...")

    from app.core.database import engine
    from app.models import Base

    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)

    await seed_pixelle_project()

    print("\n初始化完成!")


if __name__ == "__main__":
    asyncio.run(main())