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
        # For very long periods (>150 years), start from a historical date
        # to ensure we can get complete orbit data within NASA's range
        if period_days > 150 * 365:  # More than 150 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}: {e}")
            raise


orbit_service = OrbitService()