cosmo/backend/scripts/update_orbit_params.py

"""
更新太阳系行星和矮行星的轨道参数到 extra_data 字段

将硬编码在 celestial_orbit.py 中的轨道周期和颜色迁移到数据库的 extra_data 字段
这样用户可以在后台界面直接编辑这些参数
"""
import asyncio
import sys
from pathlib import Path

# Add backend directory to path
backend_dir = Path(__file__).parent.parent
sys.path.insert(0, str(backend_dir))

from sqlalchemy import select, update
from app.database import AsyncSessionLocal
from app.models.db.celestial_body import CelestialBody


# 轨道参数（从 celestial_orbit.py 迁移）
ORBIT_PARAMS = {
    # 行星 - 完整公转周期
    "199": {
        "orbit_period_days": 88.0,
        "orbit_color": "#8C7853",
        "name_zh": "水星"
    },
    "299": {
        "orbit_period_days": 224.7,
        "orbit_color": "#FFC649",
        "name_zh": "金星"
    },
    "399": {
        "orbit_period_days": 365.25,
        "orbit_color": "#4A90E2",
        "name_zh": "地球"
    },
    "499": {
        "orbit_period_days": 687.0,
        "orbit_color": "#CD5C5C",
        "name_zh": "火星"
    },
    "599": {
        "orbit_period_days": 4333.0,
        "orbit_color": "#DAA520",
        "name_zh": "木星"
    },
    "699": {
        "orbit_period_days": 10759.0,
        "orbit_color": "#F4A460",
        "name_zh": "土星"
    },
    "799": {
        "orbit_period_days": 30687.0,
        "orbit_color": "#4FD1C5",
        "name_zh": "天王星"
    },
    "899": {
        "orbit_period_days": 60190.0,
        "orbit_color": "#4169E1",
        "name_zh": "海王星"
    },

    # 矮行星 - 完整公转周期
    "999": {
        "orbit_period_days": 90560.0,
        "orbit_color": "#8B7355",
        "name_zh": "冥王星"
    },
    "2000001": {
        "orbit_period_days": 1680.0,
        "orbit_color": "#9E9E9E",
        "name_zh": "谷神星"
    },
    "136199": {
        "orbit_period_days": 203500.0,
        "orbit_color": "#E0E0E0",
        "name_zh": "阋神星"
    },
    "136108": {
        "orbit_period_days": 104000.0,
        "orbit_color": "#D4A574",
        "name_zh": "妊神星"
    },
    "136472": {
        "orbit_period_days": 112897.0,
        "orbit_color": "#C49A6C",
        "name_zh": "鸟神星"
    },
}


async def update_orbit_parameters():
    """更新数据库中的轨道参数"""

    async with AsyncSessionLocal() as session:
        print("🔄 开始更新轨道参数...\n")

        updated_count = 0
        not_found_count = 0

        for body_id, params in ORBIT_PARAMS.items():
            # 查询天体
            result = await session.execute(
                select(CelestialBody).where(CelestialBody.id == body_id)
            )
            body = result.scalar_one_or_none()

            if not body:
                print(f"⚠️  天体 {body_id} ({params['name_zh']}) 未找到")
                not_found_count += 1
                continue

            # 合并 extra_data
            extra_data = body.extra_data or {}
            extra_data["orbit_period_days"] = params["orbit_period_days"]
            extra_data["orbit_color"] = params["orbit_color"]

            # 更新数据库
            await session.execute(
                update(CelestialBody)
                .where(CelestialBody.id == body_id)
                .values(extra_data=extra_data)
            )

            print(f"✅ {params['name_zh']:8s} (ID: {body_id:7s}) - "
                  f"周期: {params['orbit_period_days']:8.1f} 天 ({params['orbit_period_days']/365.25:6.2f} 年), "
                  f"颜色: {params['orbit_color']}")
            updated_count += 1

        await session.commit()

        print(f"\n{'='*80}")
        print(f"✅ 更新完成: {updated_count} 个天体")
        if not_found_count > 0:
            print(f"⚠️  未找到: {not_found_count} 个天体")
        print(f"{'='*80}")


async def main():
    """主函数"""
    print("=" * 80)
    print("太阳系行星和矮行星轨道参数更新工具")
    print("=" * 80)
    print()

    await update_orbit_parameters()

    print("\n灶神星（Vesta）轨道参数：")
    print("=" * 80)
    print("JPL Horizons ID: 2000004")
    print("英文名: Vesta")
    print("中文名: 灶神星")
    print("类型: 矮行星 (dwarf_planet)")
    print()
    print("轨道参数:")
    print("  - 轨道周期: 1325.46 天 (约 3.63 年)")
    print("  - 建议颜色: #A8A8A8 (浅灰色)")
    print("  - 半长轴: 2.36 AU")
    print("  - 离心率: 0.089")
    print("  - 轨道倾角: 7.14°")
    print()
    print("描述: 灶神星是小行星带中第二大的小行星，直径约525公里。")
    print("      它是唯一一颗肉眼可见的小行星，也是黎明号探测器访问过的天体。")
    print("=" * 80)


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