import polars as pl


# step2 : 주문일 기간 filter
def run(tb_sales_all, options, contexts):
    month_names = [
        "January",
        "February",
        "March",
        "April",
        "May",
        "June",
        "July",
        "August",
        "September",
        "October",
        "November",
        "December",
    ]
    month_name_mapping = {i + 1: month_names[i] for i in range(len(month_names))}

    df = (
        tb_sales_all
        # 1) 문자열 → Date 변환
        .with_columns(
            pl.col("order_date").str.strptime(pl.Date, strict=False).alias("order_date")
        )
        # 2) 연/월 추출
        .with_columns(
            [
                pl.col("order_date").dt.year().alias("order_year"),
                pl.col("order_date").dt.month().alias("order_month"),
            ]
        )
        # 3) 월 이름 컬럼 추가
        .with_columns(
            [
                pl.col("order_month")
                .cast(pl.Utf8)
                .replace(month_name_mapping)  # int → str 매핑
                .alias("order_month_name")
            ]
        )
        # 4) 분기(Quarter) 생성
        .with_columns(
            [
                pl.when(pl.col("order_month").is_between(1, 3))
                .then(pl.lit("Q1"))
                .when(pl.col("order_month").is_between(4, 6))
                .then(pl.lit("Q2"))
                .when(pl.col("order_month").is_between(7, 9))
                .then(pl.lit("Q3"))
                .otherwise(pl.lit("Q4"))
                .alias("order_quarter")
            ]
        )
    )

    # 5) 연도 필터
    result = df.filter(pl.col("order_year") >= 2017)
    return {"tb_sales": result}
