cosmo_backend/app/services/horizons.py

"""
NASA JPL Horizons data query service
"""
from datetime import datetime, timedelta
from astroquery.jplhorizons import Horizons
from astropy.time import Time
import logging
import re
import httpx
import os

from app.models.celestial import Position, CelestialBody
from app.config import settings

logger = logging.getLogger(__name__)


class HorizonsService:
    """Service for querying NASA JPL Horizons system"""

    def __init__(self):
        """Initialize the service"""
        self.location = "@sun"  # Heliocentric coordinates

        # Set proxy for astroquery if configured
        # astroquery uses standard HTTP_PROXY and HTTPS_PROXY environment variables
        if settings.http_proxy:
            os.environ['HTTP_PROXY'] = settings.http_proxy
            logger.info(f"Set HTTP_PROXY for astroquery: {settings.http_proxy}")
        if settings.https_proxy:
            os.environ['HTTPS_PROXY'] = settings.https_proxy
            logger.info(f"Set HTTPS_PROXY for astroquery: {settings.https_proxy}")

    async def get_object_data_raw(self, body_id: str) -> str:
        """
        Get raw object data (terminal style text) from Horizons

        Args:
            body_id: JPL Horizons ID

        Returns:
            Raw text response from NASA
        """
        url = "https://ssd.jpl.nasa.gov/api/horizons.api"
        # Ensure ID is quoted for COMMAND
        cmd_val = f"'{body_id}'" if not body_id.startswith("'") else body_id

        params = {
            "format": "text",
            "COMMAND": cmd_val,
            "OBJ_DATA": "YES",
            "MAKE_EPHEM": "NO",
            "EPHEM_TYPE": "VECTORS",
            "CENTER": "@sun"
        }

        try:
            # Configure proxy if available
            client_kwargs = {"timeout": 5.0}
            if settings.proxy_dict:
                client_kwargs["proxies"] = settings.proxy_dict
                logger.info(f"Using proxy for NASA API: {settings.proxy_dict}")

            async with httpx.AsyncClient(**client_kwargs) as client:
                logger.info(f"Fetching raw data for body {body_id}")
                response = await client.get(url, params=params)

                if response.status_code != 200:
                    raise Exception(f"NASA API returned status {response.status_code}")

                return response.text
        except Exception as e:
            logger.error(f"Error fetching raw data for {body_id}: {str(e)}")
            raise

    def get_body_positions(
        self,
        body_id: str,
        start_time: datetime | None = None,
        end_time: datetime | None = None,
        step: str = "1d",
    ) -> list[Position]:
        """
        Get positions for a celestial body over a time range

        Args:
            body_id: JPL Horizons ID (e.g., '-31' for Voyager 1)
            start_time: Start datetime (default: now)
            end_time: End datetime (default: now)
            step: Time step (e.g., '1d' for 1 day, '1h' for 1 hour)

        Returns:
            List of Position objects
        """
        try:
            # Set default times
            if start_time is None:
                start_time = datetime.utcnow()
            if end_time is None:
                end_time = start_time

            # Convert to astropy Time objects for single point queries
            # For ranges, use ISO format strings which Horizons prefers

            # Create time range
            if start_time == end_time:
                # Single time point - use JD format
                epochs = Time(start_time).jd
            else:
                # Time range - use ISO format (YYYY-MM-DD HH:MM)
                # Horizons expects this format for ranges
                start_str = start_time.strftime('%Y-%m-%d %H:%M')
                end_str = end_time.strftime('%Y-%m-%d %H:%M')
                epochs = {"start": start_str, "stop": end_str, "step": step}

            logger.info(f"Querying Horizons for body {body_id} from {start_time} to {end_time}")

            # Query JPL Horizons
            obj = Horizons(id=body_id, location=self.location, epochs=epochs)
            vectors = obj.vectors()

            # Extract positions
            positions = []
            if isinstance(epochs, dict):
                # Multiple time points
                for i in range(len(vectors)):
                    pos = Position(
                        time=Time(vectors["datetime_jd"][i], format="jd").datetime,
                        x=float(vectors["x"][i]),
                        y=float(vectors["y"][i]),
                        z=float(vectors["z"][i]),
                    )
                    positions.append(pos)
            else:
                # Single time point
                pos = Position(
                    time=start_time,
                    x=float(vectors["x"][0]),
                    y=float(vectors["y"][0]),
                    z=float(vectors["z"][0]),
                )
                positions.append(pos)

            logger.info(f"Successfully retrieved {len(positions)} positions for body {body_id}")
            return positions

        except Exception as e:
            logger.error(f"Error querying Horizons for body {body_id}: {str(e)}")
            raise

    def search_body_by_name(self, name: str) -> dict:
        """
        Search for a celestial body by name in NASA Horizons database

        Args:
            name: Body name or ID to search for

        Returns:
            Dictionary with search results:
            {
                "success": bool,
                "id": str (extracted or input),
                "name": str (short name),
                "full_name": str (complete name from NASA),
                "error": str (if failed)
            }
        """
        try:
            logger.info(f"Searching Horizons for: {name}")

            # Try to query with the name
            obj = Horizons(id=name, location=self.location)
            vec = obj.vectors()

            # Get the full target name from response
            targetname = vec['targetname'][0]
            logger.info(f"Found target: {targetname}")

            # Extract ID and name from targetname
            # Possible formats:
            # 1. "136472 Makemake (2005 FY9)" - ID at start
            # 2. "Voyager 1 (spacecraft) (-31)" - ID in parentheses
            # 3. "Mars (499)" - ID in parentheses
            # 4. "Parker Solar Probe (spacecraft)" - no ID
            # 5. "Hubble Space Telescope (spacecra" - truncated

            numeric_id = None
            short_name = None

            # Check if input is already a numeric ID
            input_is_numeric = re.match(r'^-?\d+$', name.strip())
            if input_is_numeric:
                numeric_id = name.strip()
                # Extract name from targetname
                # Remove leading ID if present
                name_part = re.sub(r'^\d+\s+', '', targetname)
                short_name = name_part.split('(')[0].strip()
            else:
                # Try to extract ID from start of targetname (format: "136472 Makemake")
                start_match = re.match(r'^(\d+)\s+(.+)', targetname)
                if start_match:
                    numeric_id = start_match.group(1)
                    short_name = start_match.group(2).split('(')[0].strip()
                else:
                    # Try to extract ID from parentheses (format: "Name (-31)" or "Name (499)")
                    id_match = re.search(r'\((-?\d+)\)', targetname)
                    if id_match:
                        numeric_id = id_match.group(1)
                        short_name = targetname.split('(')[0].strip()
                    else:
                        # No numeric ID found, use input name as ID
                        numeric_id = name
                        short_name = targetname.split('(')[0].strip()

            return {
                "success": True,
                "id": numeric_id,
                "name": short_name,
                "full_name": targetname,
                "error": None
            }

        except Exception as e:
            error_msg = str(e)
            logger.error(f"Error searching for {name}: {error_msg}")

            # Check for specific error types
            if 'Ambiguous target name' in error_msg:
                return {
                    "success": False,
                    "id": None,
                    "name": None,
                    "full_name": None,
                    "error": "名称不唯一，请提供更具体的名称或 JPL Horizons ID"
                }
            elif 'No matches found' in error_msg or 'Unknown target' in error_msg:
                return {
                    "success": False,
                    "id": None,
                    "name": None,
                    "full_name": None,
                    "error": "未找到匹配的天体，请检查名称或 ID"
                }
            else:
                return {
                    "success": False,
                    "id": None,
                    "name": None,
                    "full_name": None,
                    "error": f"查询失败: {error_msg}"
                }


# Singleton instance
horizons_service = HorizonsService()