"""
    util/drawing.py

    Utility functions for drawing data processing
"""
from datetime import datetime, timedelta

from data.database import PowerballDraw


def get_next_draw_date(game: str | None, from_date: str | None = None) -> str | None:
    """
    Gets the next drawing date from the current or specified date for the given game
    If the current date is the next drawing date, the current date is returned

    :param game: megamillions or powerball
    :param from_date: date to get next drawing date from, if omitted the current date is used
    :return: next drawing date

    Examples:
    >>> get_next_draw_date(game='mega-millions', from_date='02/15/2026')
    '02/17/2026'
    >>> get_next_draw_date(game='powerball', from_date='02/15/2026')
    '02/16/2026'
    """
    target_days = {
        "powerball": [0, 2, 5],
        "mega-millions": [1, 4]
    }

    return_date = None
    return_date_format = '%m/%d/%Y'

    # if a valid game is not specified, return None
    if game is None or game not in target_days.keys():
        return return_date

    # if no date was specified, use today's date
    if from_date is None:
        return_date = datetime.today()
    else:
        return_date = datetime.strptime(from_date, return_date_format)

    while return_date.weekday() not in target_days[game]:
        return_date += timedelta(days=1)

    return return_date.strftime(return_date_format)


def get_powerball_object(draw_date: str, draw_result: dict) -> PowerballDraw:
    """
    :param draw_date: date for the drawing data
    :param draw_result: drawing results for the draw_date
    :return: PowerballDraw object
    
    Examples:
    >>> get_powerball_object(draw_date='2026-02-28', draw_result={'main_balls': [6, 20, 35, 54, 65], 'powerball': 10, 'multiplier': 4})
    PowerballDraw(draw_date='2026-02-28', main_ball1=6, main_ball2=20, main_ball3=35, main_ball4=54, main_ball5=65, powerball=10, power_play=4)
    """
    draw_date = datetime.strptime(draw_date, '%Y-%m-%d').date()
    mb1, mb2, mb3, mb4, mb5 = draw_result['main_balls']
    pb = draw_result['powerball']
    power_play = draw_result['multiplier']

    return PowerballDraw(draw_date=draw_date,main_ball1=mb1, main_ball2=mb2, main_ball3=mb3, main_ball4=mb4,
         main_ball5=mb5, powerball=pb, power_play=power_play)