cosmo/backend/app/services/orbit_service.py

"""
Service for managing orbital data
"""
from datetime import datetime, timedelta
from typing import List, Dict, Optional
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.dialects.postgresql import insert

from app.models.db.orbit import Orbit
from app.models.db.celestial_body import CelestialBody
from app.services.horizons import HorizonsService
import logging

logger = logging.getLogger(__name__)


class OrbitService:


    """Service for orbit CRUD operations and generation"""





    @staticmethod


    async def get_orbit(body_id: str, session: AsyncSession) -> Optional[Orbit]:


        """Get orbit data for a specific body"""


        result = await session.execute(


            select(Orbit).where(Orbit.body_id == body_id)


        )


        return result.scalar_one_or_none()





    @staticmethod


    async def get_all_orbits(


        session: AsyncSession,


        body_type: Optional[str] = None


    ) -> List[Orbit]:


        """Get all orbits, optionally filtered by body type"""


        if body_type:


            # Join with celestial_bodies to filter by type


            query = (


                select(Orbit)


                .join(CelestialBody, Orbit.body_id == CelestialBody.id)


                .where(CelestialBody.type == body_type)


            )


        else:


            query = select(Orbit)





        result = await session.execute(query)


        return list(result.scalars().all())





    @staticmethod


    async def get_all_orbits_with_bodies(


        session: AsyncSession,


        body_type: Optional[str] = None


    ) -> List[tuple[Orbit, CelestialBody]]:


        """


        Get all orbits with their associated celestial bodies in a single query.


        This is optimized to avoid N+1 query problem.





        Returns:


            List of (Orbit, CelestialBody) tuples


        """


        if body_type:


            query = (


                select(Orbit, CelestialBody)


                .join(CelestialBody, Orbit.body_id == CelestialBody.id)


                .where(CelestialBody.type == body_type)


            )


        else:


            query = (


                select(Orbit, CelestialBody)


                .join(CelestialBody, Orbit.body_id == CelestialBody.id)


            )





        result = await session.execute(query)


        return list(result.all())





    @staticmethod


    async def save_orbit(


        body_id: str,


        points: List[Dict[str, float]],


        num_points: int,


        period_days: Optional[float],


        color: Optional[str],


        session: AsyncSession


    ) -> Orbit:


        """Save or update orbit data using UPSERT"""


        stmt = insert(Orbit).values(


            body_id=body_id,


            points=points,


            num_points=num_points,


            period_days=period_days,


            color=color,


            created_at=datetime.utcnow(),


            updated_at=datetime.utcnow()


        )





        # On conflict, update all fields


        stmt = stmt.on_conflict_do_update(


            index_elements=['body_id'],


            set_={


                'points': points,


                'num_points': num_points,


                'period_days': period_days,


                'color': color,


                'updated_at': datetime.utcnow()


            }


        )





        await session.execute(stmt)


        await session.commit()





        # Fetch and return the saved orbit


        return await OrbitService.get_orbit(body_id, session)





    @staticmethod


    async def delete_orbit(body_id: str, session: AsyncSession) -> bool:


        """Delete orbit data for a specific body"""


        orbit = await OrbitService.get_orbit(body_id, session)


        if orbit:


            await session.delete(orbit)


            await session.commit()


            return True


        return False





    @staticmethod


    async def generate_orbit(


        body_id: str,


        body_name: str,


        period_days: float,


        color: Optional[str],


        session: AsyncSession,


        horizons_service: HorizonsService


    ) -> Orbit:


        """


        Generate complete orbital data for a celestial body





        Args:


            body_id: JPL Horizons ID


            body_name: Display name (for logging)


            period_days: Orbital period in days


            color: Hex color for orbit line


            session: Database session


            horizons_service: NASA Horizons API service





        Returns:


            Generated Orbit object


        """


        logger.info(f"🌌 Generating orbit for {body_name} (period: {period_days:.1f} days)")





        # Calculate number of sample points


        # Use at least 100 points for smooth ellipse


        # For very long periods, cap at 1000 to avoid excessive data


        MIN_POINTS = 100


        MAX_POINTS = 1000





        if period_days < 3650:  # < 10 years


            # For planets: aim for ~1 point per day, minimum 100


            num_points = max(MIN_POINTS, min(int(period_days), 365))


        else:  # >= 10 years


            # For outer planets and dwarf planets: monthly sampling


            num_points = min(int(period_days / 30), MAX_POINTS)





        # Calculate step size in days


        step_days = max(1, int(period_days / num_points))





        logger.info(f"  📊 Sampling {num_points} points (every {step_days} days)")





        # Query NASA Horizons for complete orbital period


        # NASA Horizons has limited date range (typically 1900-2200)


        # For very long periods, we need to limit the query range





        MAX_QUERY_YEARS = 250  # Maximum years we can query (1900-2150)


        MAX_QUERY_DAYS = MAX_QUERY_YEARS * 365





        if period_days > MAX_QUERY_DAYS:


            # For extremely long periods (>250 years), sample a partial orbit


            # Use enough data to show the orbital shape accurately


            actual_query_days = MAX_QUERY_DAYS


            start_time = datetime(1900, 1, 1)


            end_time = datetime(1900 + MAX_QUERY_YEARS, 1, 1)





            logger.warning(f"  ⚠️  Period too long ({period_days/365:.1f} years), sampling {MAX_QUERY_YEARS} years only")


            logger.info(f"  📅 Using partial orbit range: 1900-{1900 + MAX_QUERY_YEARS}")





            # Adjust sampling rate for partial orbit


            # We still want enough points to show the shape


            partial_ratio = actual_query_days / period_days


            adjusted_num_points = max(MIN_POINTS, int(num_points * 0.5))  # At least half the intended points


            step_days = max(1, int(actual_query_days / adjusted_num_points))





            logger.info(f"  📊 Adjusted sampling: {adjusted_num_points} points (every {step_days} days)")





        elif period_days > 150 * 365:  # More than 150 years but <= 250 years


            # Start from year 1900 for historical data


            start_time = datetime(1900, 1, 1)


            end_time = start_time + timedelta(days=period_days)


            logger.info(f"  📅 Using historical date range (1900-{end_time.year}) for long-period orbit")


        else:


            start_time = datetime.utcnow()


            end_time = start_time + timedelta(days=period_days)





        try:


            # Get positions from Horizons (synchronous call)


            positions = await horizons_service.get_body_positions(


                body_id=body_id,


                start_time=start_time,


                end_time=end_time,


                step=f"{step_days}d"


            )





            if not positions or len(positions) == 0:


                raise ValueError(f"No position data returned for {body_name}")





            # Convert Position objects to list of dicts


            points = [


                {"x": pos.x, "y": pos.y, "z": pos.z}


                for pos in positions


            ]





            logger.info(f"  ✅ Retrieved {len(points)} orbital points")





            # Save to database


            orbit = await OrbitService.save_orbit(


                body_id=body_id,


                points=points,


                num_points=len(points),


                period_days=period_days,


                color=color,


                session=session


            )





            logger.info(f"  💾 Saved orbit for {body_name}")


            return orbit





        except Exception as e:


            logger.error(f"  ❌ Failed to generate orbit for {body_name}: {repr(e)}")


            raise





# Singleton instance


orbit_service = OrbitService()