cosmo_backend/scripts/seed_celestial_bodies.py

#!/usr/bin/env python3
"""
Seed celestial bodies script

Adds all celestial bodies from CELESTIAL_BODIES to the database
and fetches their current positions from NASA Horizons.

Usage:
    python scripts/seed_celestial_bodies.py
"""

import sys
import os
import asyncio
from datetime import datetime

# Add backend to path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from app.database import get_db
from app.services.horizons import horizons_service
from app.services.db_service import celestial_body_service, position_service
from app.models.celestial import CELESTIAL_BODIES


async def seed_bodies():
    """Seed celestial bodies into database"""
    print("\n" + "=" * 60)
    print("🌌 Seeding Celestial Bodies")
    print("=" * 60)

    async for session in get_db():
        success_count = 0
        skip_count = 0
        error_count = 0

        total = len(CELESTIAL_BODIES)

        for idx, (body_id, info) in enumerate(CELESTIAL_BODIES.items(), 1):
            body_name = info["name"]

            try:
                # Check if body already exists
                existing_body = await celestial_body_service.get_body_by_id(body_id, session)

                if existing_body:
                    print(f"  [{idx}/{total}] ⏭️  {body_name:20s} - Already exists")
                    skip_count += 1
                    continue

                print(f"  [{idx}/{total}] 🔄 {body_name:20s} - Creating...", end='', flush=True)

                # Create body record
                body_data = {
                    "id": body_id,
                    "name": info["name"],
                    "name_zh": info.get("name_zh"),
                    "type": info["type"],
                    "description": info.get("description"),
                    "extra_data": {
                        "launch_date": info.get("launch_date"),
                        "status": info.get("status"),
                    }
                }

                await celestial_body_service.create_body(body_data, session)
                print(f" ✅ Created", flush=True)
                success_count += 1

            except Exception as e:
                print(f" ❌ Error: {str(e)}", flush=True)
                error_count += 1
                continue

        print(f"\n{'='*60}")
        print(f"📊 Summary:")
        print(f"   ✅ Created: {success_count}")
        print(f"   ⏭️  Skipped: {skip_count}")
        print(f"   ❌ Errors:  {error_count}")
        print(f"{'='*60}\n")

        break


async def sync_current_positions():
    """Fetch and store current positions for all bodies"""
    print("\n" + "=" * 60)
    print("📍 Syncing Current Positions")
    print("=" * 60)

    async for session in get_db():
        now = datetime.utcnow()
        success_count = 0
        skip_count = 0
        error_count = 0

        all_bodies = await celestial_body_service.get_all_bodies(session)
        total = len(all_bodies)

        for idx, body in enumerate(all_bodies, 1):
            body_id = body.id
            body_name = body.name

            try:
                # Check if we have recent position (within last hour)
                from datetime import timedelta
                recent_time = now - timedelta(hours=1)
                existing_positions = await position_service.get_positions(
                    body_id, recent_time, now, session
                )

                if existing_positions and len(existing_positions) > 0:
                    print(f"  [{idx}/{total}] ⏭️  {body_name:20s} - Recent data exists")
                    skip_count += 1
                    continue

                print(f"  [{idx}/{total}] 🔄 {body_name:20s} - Fetching...", end='', flush=True)

                # Special handling for Sun
                if body_id == "10":
                    positions_data = [{"time": now, "x": 0.0, "y": 0.0, "z": 0.0}]
                # Special handling for Cassini
                elif body_id == "-82":
                    cassini_date = datetime(2017, 9, 15, 11, 58, 0)
                    positions_data = horizons_service.get_body_positions(
                        body_id, cassini_date, cassini_date
                    )
                    positions_data = [
                        {"time": p.time, "x": p.x, "y": p.y, "z": p.z}
                        for p in positions_data
                    ]
                else:
                    # Query current position
                    positions_data = horizons_service.get_body_positions(
                        body_id, now, now
                    )
                    positions_data = [
                        {"time": p.time, "x": p.x, "y": p.y, "z": p.z}
                        for p in positions_data
                    ]

                # Store positions
                for pos_data in positions_data:
                    await position_service.save_position(
                        body_id=body_id,
                        time=pos_data["time"],
                        x=pos_data["x"],
                        y=pos_data["y"],
                        z=pos_data["z"],
                        source="nasa_horizons",
                        session=session,
                    )

                print(f" ✅ Saved {len(positions_data)} position(s)", flush=True)
                success_count += 1

                # Small delay to avoid overwhelming NASA API
                await asyncio.sleep(0.5)

            except Exception as e:
                print(f" ❌ Error: {str(e)}", flush=True)
                error_count += 1
                continue

        print(f"\n{'='*60}")
        print(f"📊 Summary:")
        print(f"   ✅ Success: {success_count}")
        print(f"   ⏭️  Skipped: {skip_count}")
        print(f"   ❌ Errors:  {error_count}")
        print(f"{'='*60}\n")

        break


async def main():
    print("\n🚀 Celestial Bodies Database Seeding")
    print("=" * 60)
    print("This script will:")
    print("  1. Add all celestial bodies to the database")
    print("  2. Fetch and store their current positions")
    print("=" * 60)

    # Seed celestial bodies
    await seed_bodies()

    # Sync current positions
    await sync_current_positions()

    print("\n🎉 Seeding complete!")


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