insights/products/surveys/backend/util.py

from enum import StrEnum
from uuid import UUID

from insights.insightsql.escape_sql import escape_datastore_string

from insights.models.event.new_events_schema import events_read_table
from insights.models.property.util import get_property_string_expr

from products.surveys.backend.models import SurveyResponseArchive


class SurveyEventName(StrEnum):
    SHOWN = "survey shown"
    DISMISSED = "survey dismissed"
    SENT = "survey sent"


class SurveyEventProperties(StrEnum):
    SURVEY_ID = "$survey_id"
    SURVEY_RESPONSE = "$survey_response"
    SURVEY_ITERATION = "$survey_iteration"
    SURVEY_PARTIALLY_COMPLETED = "$survey_partially_completed"
    SURVEY_SUBMISSION_ID = "$survey_submission_id"
    SURVEY_RESPONDED = "$survey_responded"
    SURVEY_DISMISSED = "$survey_dismissed"
    SURVEY_COMPLETED = "$survey_completed"
    SURVEY_LAST_SEEN_DATE = "$last_seen_survey_date"


def get_survey_property_string_expr(
    property_name: SurveyEventProperties,
    table_alias: str | None = None,
    use_new_events_schema: bool = False,
) -> str:
    expression, _ = get_property_string_expr(
        "events",
        str(property_name),
        escape_datastore_string(str(property_name)),
        "properties",
        table_alias=table_alias,
        use_new_events_schema=use_new_events_schema,
    )
    return expression


def get_survey_property_bool_expr(property_name: SurveyEventProperties, table_alias: str | None = None) -> str:
    expression = get_survey_property_string_expr(property_name, table_alias)
    normalized_expression = f"toString(nullIf(nullIf({expression}, ''), 'null'))"
    return f"toBool(transform({normalized_expression}, ['true', 'false'], [1, 0], NULL))"


def get_survey_response_datastore_query(
    question_index: int, question_id: str | None = None, is_multiple_choice: bool = False
) -> str:
    """
    Generate a Datastore query to extract survey response based on question index or ID

    Args:
        question_index: Index of the question
        question_id: ID of the question (optional)

    Returns:
        The survey response or empty string if not found
    """
    _DANGEROUS_id_based_key = _build_id_based_key(question_index, question_id)
    index_based_key = _build_index_based_key(question_index)

    if is_multiple_choice is True:
        return _build_multiple_choice_query(_DANGEROUS_id_based_key, index_based_key)

    return _build_coalesce_query(_DANGEROUS_id_based_key, index_based_key)


def _build_id_based_key(question_index: int, question_id: str | None = None) -> str:
    if question_id:
        return escape_datastore_string(f"{SurveyEventProperties.SURVEY_RESPONSE}_{question_id}")

    # Extract the ID from the question at the given index in the questions array
    return f"CONCAT({escape_datastore_string(SurveyEventProperties.SURVEY_RESPONSE + '_')}, JSONExtractString(JSONExtractArrayRaw(properties, '$survey_questions')[{int(question_index) + 1}], 'id'))"


def _build_index_based_key(question_index: int) -> str:
    if question_index == 0:
        return f"{SurveyEventProperties.SURVEY_RESPONSE}"
    return f"{SurveyEventProperties.SURVEY_RESPONSE}_{question_index}"


def _build_coalesce_query(_DANGEROUS_id_based_key: str, index_based_key: str) -> str:
    return f"""COALESCE(
        NULLIF(JSONExtractString(properties, {_DANGEROUS_id_based_key}), ''),
        NULLIF(JSONExtractString(properties, {escape_datastore_string(index_based_key)}), '')
    )"""


def _build_multiple_choice_query(_DANGEROUS_id_based_key: str, index_based_key: str) -> str:
    return f"""if(
        JSONHas(properties, {_DANGEROUS_id_based_key}) AND length(JSONExtractArrayRaw(properties, {_DANGEROUS_id_based_key})) > 0,
        JSONExtractArrayRaw(properties, {_DANGEROUS_id_based_key}),
        JSONExtractArrayRaw(properties, {escape_datastore_string(index_based_key)})
    )"""


def get_unique_survey_event_uuids_sql_subquery(
    base_conditions_sql: list[str],
    group_by_prefix_expressions: list[str] | None = None,
    use_new_events_schema: bool = False,
) -> str:
    """
    Generates a SQL subquery string that returns unique event UUIDs for 'survey sent' events,
    deduplicating based on $survey_submission_id (for new events) or uuid (for older events) using argMax.

    The subquery is intended to be used in a `WHERE uuid IN (...)` clause.

    Args:
        base_conditions_sql: A list of SQL conditions for the WHERE clause of the subquery.
                             Example: ["team_id = %(team_id)s", "timestamp >= '2023-01-01'", "event = 'survey sent'"]
                             Callers should ensure "event = 'survey sent'" is included if that's the target.
        group_by_prefix_expressions: A list of SQL expressions to prefix the GROUP BY clause.
                                     These define the segments within which deduplication occurs.
                                     Example: ['team_id', "replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(properties, '$survey_id'), ''), 'null'), '^\"|\"$', '')"]
                                     If empty, deduplication is based purely on submission ID / UUID across
                                     all events matching base_conditions_sql.

    Returns:
        A string for the SQL subquery, e.g.,
        "(SELECT argMax(uuid, timestamp) FROM events WHERE ... GROUP BY ...)"
    """
    if not base_conditions_sql:
        raise ValueError("base_conditions_sql cannot be empty. Provide at least one condition.")

    if group_by_prefix_expressions is None:
        group_by_prefix_expressions = []

    # Ensure the event filter is present in the base conditions
    if base_conditions_sql.count(f"event = '{SurveyEventName.SENT}'") == 0:
        sql_conditions = [*base_conditions_sql, f"event = '{SurveyEventName.SENT}'"]
    else:
        sql_conditions = base_conditions_sql

    # Always include the survey_id in the group by
    survey_id_expr = get_survey_property_string_expr(
        SurveyEventProperties.SURVEY_ID, use_new_events_schema=use_new_events_schema
    )
    if survey_id_expr not in group_by_prefix_expressions:
        group_by_prefix_expressions.append(survey_id_expr)

    where_clause = " AND ".join(sql_conditions)

    submission_id_col = get_survey_property_string_expr(
        SurveyEventProperties.SURVEY_SUBMISSION_ID, use_new_events_schema=use_new_events_schema
    )
    deduplication_group_by_key = (
        f"CASE WHEN COALESCE({submission_id_col}, '') = '' THEN toString(uuid) ELSE {submission_id_col} END"
    )

    group_by_clause = ", ".join([*group_by_prefix_expressions, deduplication_group_by_key])

    events_table = events_read_table(use_new_events_schema)
    return f"(SELECT argMax(uuid, timestamp) FROM {events_table} WHERE {where_clause} GROUP BY {group_by_clause})"


def get_archived_response_uuids(survey_id: str | UUID | None, team_id: int) -> set[str]:
    """
    Get set of archived response UUIDs from PostgreSQL.

    Args:
        survey_id: UUID of the survey, or None to get all archived responses for the team
        team_id: Team ID

    Returns:
        Set of archived response UUIDs as strings
    """
    filter_kwargs: dict[str, int | str | UUID] = {"team_id": team_id}
    if survey_id is not None:
        filter_kwargs["survey_id"] = survey_id

    return {
        str(uuid)
        for uuid in SurveyResponseArchive.objects.filter(**filter_kwargs).values_list("response_uuid", flat=True)
    }