src/url_state.py

"""
URL state encoding/decoding for shareable search URLs.

Core principle: All searches are represented uniformly - basic search is just
advanced search with a single text query at weight +1.
"""

import base64
import json
from typing import Dict, List, Optional, Tuple, Any

try:
    from flask import request
except ImportError:
    request = None


def encode_search_state(
    text_queries: List[str],
    text_weights: List[float],
    image_queries: List[Tuple[float, float, float]],  # (ra, dec, fov)
    image_weights: List[float],
    rmag_min: float = 13.0,
    rmag_max: float = 20.0,
    expand_galaxy: Optional[str] = None
) -> str:
    """
    Encode search state into a compact base64 string.

    Args:
        text_queries: List of text query strings
        text_weights: List of weights for text queries
        image_queries: List of (ra, dec, fov) tuples
        image_weights: List of weights for image queries
        rmag_min: Minimum r-band magnitude (default 13.0)
        rmag_max: Maximum r-band magnitude (default 20.0)
        expand_galaxy: Optional object_id to expand in modal after search

    Returns:
        URL-safe base64 encoded string
    """
    state = {}

    # Only include non-empty query arrays
    if text_queries:
        state["tq"] = text_queries
        state["tw"] = text_weights

    if image_queries:
        state["iq"] = image_queries
        state["iw"] = image_weights

    # Only include magnitude limits if they differ from defaults
    if rmag_min != 13.0:
        state["rmin"] = rmag_min
    if rmag_max != 20.0:
        state["rmax"] = rmag_max

    # Include expand galaxy if specified
    if expand_galaxy:
        state["exp"] = expand_galaxy

    # Encode to JSON then base64 (URL-safe)
    json_str = json.dumps(state, separators=(',', ':'))  # Compact JSON
    encoded = base64.urlsafe_b64encode(json_str.encode('utf-8')).decode('utf-8')

    # Remove padding (= characters) to make URL shorter - we'll add back on decode
    return encoded.rstrip('=')


def decode_search_state(encoded: str) -> Dict[str, Any]:
    """
    Decode a base64-encoded search state string.

    Args:
        encoded: URL-safe base64 encoded string

    Returns:
        Dictionary with full keys:
        {
            'text_queries': [...],
            'text_weights': [...],
            'image_queries': [...],
            'image_weights': [...],
            'rmag_min': float,
            'rmag_max': float,
            'expand_galaxy': str or None
        }

    Returns empty/default state if decoding fails.
    """
    # Default/empty state
    default_state = {
        'text_queries': [],
        'text_weights': [],
        'image_queries': [],
        'image_weights': [],
        'rmag_min': 13.0,
        'rmag_max': 20.0,
        'expand_galaxy': None
    }

    if not encoded:
        return default_state

    try:
        # Add back padding if needed
        padding = 4 - (len(encoded) % 4)
        if padding != 4:
            encoded += '=' * padding

        # Decode from base64 then JSON
        json_str = base64.urlsafe_b64decode(encoded.encode('utf-8')).decode('utf-8')
        compact_state = json.loads(json_str)

        # Expand compact keys to full keys
        state = {
            'text_queries': compact_state.get('tq', []),
            'text_weights': compact_state.get('tw', []),
            'image_queries': compact_state.get('iq', []),
            'image_weights': compact_state.get('iw', []),
            'rmag_min': compact_state.get('rmin', 13.0),
            'rmag_max': compact_state.get('rmax', 20.0),
            'expand_galaxy': compact_state.get('exp', None)
        }

        return state

    except (ValueError, KeyError, json.JSONDecodeError, Exception):
        # Silently fail and return empty state for any decode errors
        return default_state


def build_share_url(
    text_queries: List[str],
    text_weights: List[float],
    image_queries: List[Tuple[float, float, float]],
    image_weights: List[float],
    rmag_min: float = 13.0,
    rmag_max: float = 20.0,
    expand_galaxy: Optional[str] = None
) -> str:
    """
    Build a complete shareable URL for the current search state.

    Args:
        Same as encode_search_state()

    Returns:
        Complete URL like: https://app.com?s=<encoded_state>
    """
    from src.config import URL_STATE_PARAM

    # Get base URL from Flask request context
    try:
        if request and hasattr(request, 'host_url'):
            base_url = request.host_url.rstrip('/')
        else:
            base_url = "http://localhost:8050"
    except (RuntimeError, AttributeError):
        # If not in request context (e.g., during testing), use a placeholder
        base_url = "http://localhost:8050"

    # Encode the state
    encoded = encode_search_state(
        text_queries=text_queries,
        text_weights=text_weights,
        image_queries=image_queries,
        image_weights=image_weights,
        rmag_min=rmag_min,
        rmag_max=rmag_max,
        expand_galaxy=expand_galaxy
    )

    return f"{base_url}?{URL_STATE_PARAM}={encoded}"


def parse_url_search_param(search_string: str) -> Dict[str, Any]:
    """
    Parse the search parameter from a URL query string.

    Args:
        search_string: The URL search/query string (e.g., "?s=eyJxdWVy...")

    Returns:
        Decoded state dictionary (same format as decode_search_state)
    """
    from src.config import URL_STATE_PARAM

    if not search_string:
        return decode_search_state("")

    # Remove leading '?' if present
    search_string = search_string.lstrip('?')

    # Parse query parameters
    params = {}
    for param in search_string.split('&'):
        if '=' in param:
            key, value = param.split('=', 1)
            params[key] = value

    # Get the state parameter
    encoded = params.get(URL_STATE_PARAM, '')

    return decode_search_state(encoded)