src/validation.py

"""Validate client state before rendering results or issuing upstream queries."""

import math

from src.config import DEFAULT_TOP_K, ZILLIZ_PRIMARY_KEY
from src.search_limits import MAX_SEARCH_TERMS, SearchLimitError, validate_search_size

# Includes result arrays and dynamic components, not just the search terms.
MAX_REQUEST_BYTES = 2 * 1024 * 1024
MAX_ENCODED_STATE_CHARS = 256 * 1024


def finite_number(value, name, lower=None, upper=None):
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise SearchLimitError(f"{name} must be a finite number.")
    try:
        valid = math.isfinite(value)
    except OverflowError:
        valid = False
    if not valid:
        raise SearchLimitError(f"{name} must be a finite number.")
    if lower is not None and value < lower or upper is not None and value > upper:
        raise SearchLimitError(f"{name} must be between {lower} and {upper}.")
    return value


def validate_magnitudes(rmag_min, rmag_max):
    for value in (rmag_min, rmag_max):
        if value is not None:
            finite_number(value, "Magnitude", 13, 20)
    if rmag_min is not None and rmag_max is not None and rmag_min > rmag_max:
        raise SearchLimitError("The minimum magnitude must not exceed the maximum.")


def parse_magnitude_range(values):
    if values is None:
        return None, None
    if not isinstance(values, (list, tuple)) or len(values) != 2:
        raise SearchLimitError("Supply a minimum and maximum magnitude.")
    validate_magnitudes(*values)
    return values[0], values[1]


def validate_image_query(image):
    if not isinstance(image, dict) or 'ra' not in image or 'dec' not in image:
        raise SearchLimitError("Each image query needs RA and Dec coordinates.")
    finite_number(image['ra'], "RA", 0, 360)
    finite_number(image['dec'], "Dec", -90, 90)
    fov = finite_number(image.get('fov', 0.025), "Field of view", 0, 180)
    if fov == 0:
        raise SearchLimitError("Field of view must be greater than zero.")


def validate_weights(queries, weights):
    queries = [] if queries is None else queries
    weights = [] if weights is None else weights
    if not isinstance(weights, (list, tuple)) or len(weights) != len(queries):
        raise SearchLimitError("Each search term must have exactly one weight.")
    for weight in weights:
        finite_number(weight, "Weight", -10, 10)


def validate_search_request(text_queries=None, text_weights=None, image_queries=None,
                            image_weights=None, rmag_min=None, rmag_max=None,
                            top_k=DEFAULT_TOP_K):
    validate_search_size(text_queries, image_queries)
    validate_weights(text_queries, text_weights)
    validate_weights(image_queries, image_weights)
    validate_magnitudes(rmag_min, rmag_max)
    if type(top_k) is not int or not 1 <= top_k <= DEFAULT_TOP_K:
        raise SearchLimitError(f"Request between 1 and {DEFAULT_TOP_K} results.")
    for image in image_queries or []:
        validate_image_query(image)


def parse_advanced_queries(query, operation, query_type, ra, dec,
                           query_types, texts, ras, decs, operations):
    arrays = (query_types, texts, ras, decs, operations)
    if any(not isinstance(values, (list, tuple)) for values in arrays):
        raise SearchLimitError("Additional query fields must be lists.")
    if len({len(values) for values in arrays}) != 1:
        raise SearchLimitError("Additional query fields must have matching lengths.")
    if len(query_types) > MAX_SEARCH_TERMS:
        raise SearchLimitError(f"Use at most {MAX_SEARCH_TERMS} text and image terms per search.")

    text_queries, text_weights, image_queries, image_weights = [], [], [], []
    rows = [(query_type, query, ra, dec, operation)]
    rows.extend(zip(query_types, texts, ras, decs, operations))
    for kind, text, ra, dec, op in rows:
        if kind not in ('text', 'image'):
            raise SearchLimitError("Query type must be text or image.")
        try:
            weight = {'+': 1.0, '-': -1.0}[op] if op in ('+', '-') else float(op)
        except (TypeError, ValueError, OverflowError):
            raise SearchLimitError("Each query needs a numeric weight.") from None
        if isinstance(op, bool):
            raise SearchLimitError("Each query needs a numeric weight.")
        finite_number(weight, "Weight", -10, 10)
        if kind == 'text':
            if text is not None and not isinstance(text, str):
                raise SearchLimitError("Each text term must be text.")
            if text and text.strip():
                text_queries.append(text.strip())
                text_weights.append(weight)
        elif ra is not None or dec is not None:
            image = {'ra': ra, 'dec': dec, 'fov': 0.025}
            validate_image_query(image)
            image_queries.append(image)
            image_weights.append(weight)
    validate_search_request(text_queries, text_weights, image_queries, image_weights)
    return text_queries, text_weights, image_queries, image_weights


def validate_result_state(data):
    """Result stores are untrusted even though they originated on the server."""
    if not isinstance(data, dict):
        raise SearchLimitError("Invalid results. Please run the search again.")
    columns = (ZILLIZ_PRIMARY_KEY, 'ra', 'dec', 'r_mag', 'distance', 'cutout_url')
    if any(not isinstance(data.get(key), list) for key in columns):
        raise SearchLimitError("Invalid result columns. Please run the search again.")
    count = len(data['ra'])
    if count > DEFAULT_TOP_K or any(len(data[key]) != count for key in columns):
        raise SearchLimitError(f"Results must have matching columns and at most {DEFAULT_TOP_K} galaxies.")
    loaded = data.get('loaded_count')
    if type(loaded) is not int or not 0 <= loaded <= count:
        raise SearchLimitError("Invalid pagination. Please run the search again.")
    if not isinstance(data.get('query'), str) or len(data['query']) > 20_000:
        raise SearchLimitError("Invalid result description. Please run the search again.")
    for index in range(count):
        finite_number(data['ra'][index], "RA", 0, 360)
        finite_number(data['dec'][index], "Dec", -90, 90)
        finite_number(data['r_mag'][index], "Magnitude")
        finite_number(data['distance'][index], "Similarity")
        pk, url = data[ZILLIZ_PRIMARY_KEY][index], data['cutout_url'][index]
        if isinstance(pk, bool) or not isinstance(pk, (str, int)) or len(str(pk)) > 512:
            raise SearchLimitError("Invalid galaxy identifier.")
        if not isinstance(url, str) or len(url) > 2048:
            raise SearchLimitError("Invalid galaxy image URL.")
    return count