"""
Helper functions for data endpoints.
"""

from datetime import datetime
from typing import Any, Callable, Optional, Type, TypeVar, cast

from pydantic import BaseModel

from ....connectors import snowflake_db
from ....typings import SqlQuery
from ....utils.generic import lowercase_keys
from . import typings

T = TypeVar("T", bound=BaseModel)


def args_extract_label_ids(label_ids: Optional[str]) -> list[typings.LabelId] | None:
    """
    Extract label IDs from the provided arguments.
    The label IDs are expected to be in a comma-separated format.

    The function returns a list of unique integers representing the label IDs,
    sorted in ascending order. If the input is empty or None, it returns None.
    """
    if label_ids is None or not label_ids.strip():
        return None
    unique_label_ids = set(label_ids.split(","))
    sorted_label_ids = sorted(
        int(x.strip()) for x in unique_label_ids if x.strip().isdigit()
    )
    if not sorted_label_ids:
        raise ValueError("No valid label IDs provided.")
    return sorted_label_ids


async def get_formatted_data(
    snowflake_client: snowflake_db.Client,
    *,
    sql_query: SqlQuery,
    model_class: Type[T],
    hook_after_fetch: Optional[Callable[[list[dict[str, Any]]], None]] = None,
) -> list[T]:
    """
    Fetch and format data from Snowflake using the provided SQL query and model class.
    This function combines fetching and formatting into a single step.

    An optional hook can be provided to modify the raw data in place after fetching.
    """
    raw_data = await _fetch_data(snowflake_client, sql_query)
    if hook_after_fetch:
        hook_after_fetch(raw_data)
    return [model_class(**row) for row in raw_data]


async def _fetch_data(snowflake_client: snowflake_db.Client, sql_query: SqlQuery):
    """
    Fetch data from Snowflake using the provided SQL query.
    This function is a wrapper around the Snowflake client fetch method.
    """
    raw_data = await snowflake_client.afetch_all(sql_query)
    raw_data = cast(list[dict[str, Any]], raw_data)
    return lowercase_keys(raw_data)


def validate_date_yearmonth(date_yearmonth: str) -> None:
    """
    Validate the format and validity of a date in 'YYYY-MM' format.

    Raises ValueError if the format is invalid or if the date is not valid.
    """
    try:
        datetime.strptime(date_yearmonth, "%Y-%m")
    except ValueError as ex:
        raise ValueError(
            f"Invalid date_yearmonth format: {date_yearmonth}. "
            "Expected format is 'YYYY-MM'."
        ) from ex