airlabs-manage/backend/seed.py

"""种子数据 —— 用于测试和演示"""
from datetime import date, timedelta
from database import SessionLocal, engine, Base
from models import *
from auth import hash_password

Base.metadata.create_all(bind=engine)
db = SessionLocal()


def seed():
    # 清空数据
    for table in reversed(Base.metadata.sorted_tables):
        db.execute(table.delete())
    db.commit()

    # ── 用户 ──
    users = [
        User(username="admin", password_hash=hash_password("admin123"),
             name="老板", phase_group=PhaseGroup.PRODUCTION, role=UserRole.OWNER, monthly_salary=0),
        User(username="zhangsan", password_hash=hash_password("123456"),
             name="张三", phase_group=PhaseGroup.PRE, role=UserRole.LEADER, monthly_salary=15000),
        User(username="lisi", password_hash=hash_password("123456"),
             name="李四", phase_group=PhaseGroup.PRODUCTION, role=UserRole.LEADER, monthly_salary=18000),
        User(username="wangwu", password_hash=hash_password("123456"),
             name="王五", phase_group=PhaseGroup.PRODUCTION, role=UserRole.MEMBER, monthly_salary=12000),
        User(username="zhaoliu", password_hash=hash_password("123456"),
             name="赵六", phase_group=PhaseGroup.PRODUCTION, role=UserRole.MEMBER, monthly_salary=12000),
        User(username="sunqi", password_hash=hash_password("123456"),
             name="孙七", phase_group=PhaseGroup.POST, role=UserRole.MEMBER, monthly_salary=13000),
        User(username="producer", password_hash=hash_password("123456"),
             name="陈制片", phase_group=PhaseGroup.PRODUCTION, role=UserRole.SUPERVISOR, monthly_salary=20000),
    ]
    db.add_all(users)
    db.flush()

    admin, zhangsan, lisi, wangwu, zhaoliu, sunqi, producer = users

    # ── 项目 ──
    proj_a = Project(
        name="星际漫游 第一季", project_type=ProjectType.CLIENT_FORMAL,
        leader_id=lisi.id, current_phase=PhaseGroup.PRODUCTION,
        episode_duration_minutes=5, episode_count=13,
        estimated_completion_date=date.today() + timedelta(days=60),
        contract_amount=100000,
    )
    proj_b = Project(
        name="品牌方 TVC", project_type=ProjectType.CLIENT_FORMAL,
        leader_id=lisi.id, current_phase=PhaseGroup.PRODUCTION,
        episode_duration_minutes=1, episode_count=3,
        estimated_completion_date=date.today() + timedelta(days=20),
        contract_amount=50000,
    )
    proj_c = Project(
        name="甲方风格测试", project_type=ProjectType.CLIENT_TEST,
        leader_id=zhangsan.id, current_phase=PhaseGroup.PRE,
        episode_duration_minutes=1, episode_count=1,
    )
    proj_d = Project(
        name="AI 短剧原创", project_type=ProjectType.INTERNAL_ORIGINAL,
        leader_id=lisi.id, current_phase=PhaseGroup.PRE,
        episode_duration_minutes=8, episode_count=6,
        estimated_completion_date=date.today() + timedelta(days=90),
    )
    db.add_all([proj_a, proj_b, proj_c, proj_d])
    db.flush()

    # ── 内容提交（模拟近两周的数据） ──
    base_date = date.today() - timedelta(days=14)
    submissions = []

    # 张三（前期组）给项目 A 和 D 做前期
    for i in range(5):
        d = base_date + timedelta(days=i)
        submissions.append(Submission(
            user_id=zhangsan.id, project_id=proj_a.id,
            project_phase=PhaseGroup.PRE, work_type=WorkType.PLAN,
            content_type=ContentType.DESIGN, total_seconds=0,
            submit_to=SubmitTo.INTERNAL, description=f"角色设定第{i+1}版",
            submit_date=d,
        ))
    for i in range(3):
        d = base_date + timedelta(days=i + 5)
        submissions.append(Submission(
            user_id=zhangsan.id, project_id=proj_d.id,
            project_phase=PhaseGroup.PRE, work_type=WorkType.PLAN,
            content_type=ContentType.DESIGN, total_seconds=0,
            submit_to=SubmitTo.INTERNAL, description=f"剧本大纲第{i+1}稿",
            submit_date=d,
        ))

    # 李四（制作组组长）主要做项目 A
    for i in range(10):
        d = base_date + timedelta(days=i)
        secs = 45 + (i % 3) * 15  # 45-75秒
        wt = WorkType.TEST if i < 2 else WorkType.PRODUCTION
        submissions.append(Submission(
            user_id=lisi.id, project_id=proj_a.id,
            project_phase=PhaseGroup.PRODUCTION, work_type=wt,
            content_type=ContentType.ANIMATION, total_seconds=secs,
            duration_minutes=secs // 60, duration_seconds=secs % 60,
            submit_to=SubmitTo.INTERNAL, description=f"第1集片段 - 场景{i+1}",
            submit_date=d,
        ))

    # 王五（制作组）做项目 A 和 B
    for i in range(8):
        d = base_date + timedelta(days=i)
        secs = 30 + (i % 4) * 20  # 30-90秒
        submissions.append(Submission(
            user_id=wangwu.id, project_id=proj_a.id,
            project_phase=PhaseGroup.PRODUCTION, work_type=WorkType.PRODUCTION,
            content_type=ContentType.ANIMATION, total_seconds=secs,
            duration_minutes=secs // 60, duration_seconds=secs % 60,
            submit_to=SubmitTo.LEADER, description=f"第2集动画片段{i+1}",
            submit_date=d,
        ))
    for i in range(4):
        d = base_date + timedelta(days=i + 8)
        secs = 20 + i * 10
        submissions.append(Submission(
            user_id=wangwu.id, project_id=proj_b.id,
            project_phase=PhaseGroup.PRODUCTION, work_type=WorkType.PRODUCTION,
            content_type=ContentType.ANIMATION, total_seconds=secs,
            duration_minutes=secs // 60, duration_seconds=secs % 60,
            submit_to=SubmitTo.LEADER, description=f"TVC 片段{i+1}",
            submit_date=d,
        ))

    # 赵六（制作组）做项目 A
    for i in range(10):
        d = base_date + timedelta(days=i)
        secs = 50 + (i % 2) * 30  # 50-80秒
        wt = WorkType.TEST if i < 1 else WorkType.PRODUCTION
        submissions.append(Submission(
            user_id=zhaoliu.id, project_id=proj_a.id,
            project_phase=PhaseGroup.PRODUCTION, work_type=wt,
            content_type=ContentType.ANIMATION, total_seconds=secs,
            duration_minutes=secs // 60, duration_seconds=secs % 60,
            submit_to=SubmitTo.LEADER, description=f"第3集场景动画{i+1}",
            submit_date=d,
        ))

    # 孙七（后期组）剪辑
    for i in range(3):
        d = base_date + timedelta(days=i + 10)
        submissions.append(Submission(
            user_id=sunqi.id, project_id=proj_a.id,
            project_phase=PhaseGroup.POST, work_type=WorkType.PRODUCTION,
            content_type=ContentType.EDITING, total_seconds=0,
            submit_to=SubmitTo.PRODUCER, description=f"第{i+1}集粗剪完成",
            submit_date=d,
        ))
    # 后期补拍
    submissions.append(Submission(
        user_id=sunqi.id, project_id=proj_a.id,
        project_phase=PhaseGroup.POST, work_type=WorkType.PRODUCTION,
        content_type=ContentType.ANIMATION, total_seconds=15,
        duration_seconds=15,
        submit_to=SubmitTo.PRODUCER, description="第1集补拍修改镜头",
        submit_date=base_date + timedelta(days=12),
    ))

    db.add_all(submissions)

    # ── AI 工具成本 ──
    db.add(AIToolCost(
        tool_name="Midjourney", subscription_period=SubscriptionPeriod.MONTHLY,
        amount=200, allocation_type=CostAllocationType.TEAM,
        recorded_by=producer.id, record_date=date.today().replace(day=1),
    ))
    db.add(AIToolCost(
        tool_name="Runway", subscription_period=SubscriptionPeriod.MONTHLY,
        amount=600, allocation_type=CostAllocationType.PROJECT,
        project_id=proj_a.id,
        recorded_by=producer.id, record_date=date.today().replace(day=1),
    ))

    # ── 外包成本 ──
    db.add(OutsourceCost(
        project_id=proj_a.id, outsource_type=OutsourceType.ANIMATION,
        episode_start=10, episode_end=13, amount=20000,
        recorded_by=producer.id, record_date=date.today() - timedelta(days=5),
    ))

    db.commit()
    print("[OK] seed data generated")
    print(f"  - users: {len(users)}")
    print(f"  - projects: 4")
    print(f"  - submissions: {len(submissions)}")
    print(f"  - default account: admin / admin123")


if __name__ == "__main__":
    seed()