import polars as pl


def run(retail_inventory, competitor_prices, options, contexts):
    """
    Updates competitor price information in the retail data using competitor price updates.

    Parameters:
        retail_inventory: Original retail information DataFrame
        competitor_prices: Competitor price update information DataFrame
        options: Additional options (not used in this implementation)

    Returns:
        Dictionary with 'update_competitor_price' key containing the updated retail DataFrame
    """
    # Derive new_competitor_price from price_krw
    if competitor_prices.schema["price_krw"] == pl.String:
        # CSV source: "4,660 ₩" → strip currency symbol/commas, then cast to int
        competitor_prices = competitor_prices.with_columns(
            pl.col("price_krw")
            .str.replace_all(" ₩", "")
            .str.replace_all(",", "")
            .cast(pl.Int64)
            .alias("new_competitor_price")
        )
    else:
        # Parquet source: already numeric
        competitor_prices = competitor_prices.with_columns(
            pl.col("price_krw").cast(pl.Int64).alias("new_competitor_price")
        )

    # Select only item_code and new price, deduplicate by item_code (take first occurrence)
    price_updates = competitor_prices.select(
        ["item_code", "new_competitor_price"]
    ).unique(subset=["item_code"], keep="first")

    # Left join retail with competitor prices on item_code only
    result = retail_inventory.join(price_updates, on=["item_code"], how="left")

    # Update competitor_price_krw_synth with new price where available
    result = result.with_columns(
        pl.when(pl.col("new_competitor_price").is_not_null())
        .then(pl.col("new_competitor_price"))
        .otherwise(pl.col("competitor_price_krw_synth"))
        .alias("competitor_price_krw_synth")
    ).drop("new_competitor_price")

    return {"update_competitor_price": result}


if __name__ == "__main__":
    # Sample retail data
    retail = pl.DataFrame(
        {
            "id": ["GNM0_7224000894", "NJ10_7224000894", "GNM0_7224000894"],
            "item_code": ["7224000894", "7224000894", "7224000894"],
            "product_name": ["이니스프리_노세범_미네랄_파우더"] * 3,
            "branch_code": ["GNM0", "NJ10", "GNM0"],
            "work_order_no": ["202408001043", "202408002098", "202408001043"],
            "PROD_GRP": ["YT44402461", "YT44402360", "YT44402606"],
            "branch_name": ["강남파이낸스점", "나주빛가람동점", "강남파이낸스점"],
            "warehouse_location": ["WH_GNM_001", "WH_NJ_001", "WH_GNM_002"],
            "warn_clearnace": [True, True, True],
            "delivery_date": ["2024-09-04"] * 3,
            "delivery_time": ["20:24:00", "16:39:00", "14:55:00"],
            "manager_id": ["HV79230589", "HV79240312", "HV79230589"],
            "inventory_qty": [51, 52, 65],
            "outbound_qty": [31, 32, 39],
            "stock_variance": [20, 20, 26],
            "inventory_turnover_rate": [25, 31, 39],
            "safety_stock_level": [15, 12, 19],
            "list_price_krw_synth": [9000, 9000, 9000],
            "selling_price_krw_synth": [4490, 5180, 5400],
            "discount_rate_synth": [0.5011, 0.4244, 0.4],
            "unit_cost_krw_synth": [2120, 2340, 2130],
            "promo_flag_synth": [1, 1, 1],
            "competitor_price_krw_synth": [4500, 5180, 5670],
            "stockout_risk_flag_synth": [0, 0, 0],
            "latitude": [37.5002, 35.0154, 37.5002],
            "longitude": [127.0365, 126.8007, 127.0365],
            "region_name": ["서울 강남", "전남 나주", "서울 강남"],
            "supplier_name": ["이니스프리"] * 3,
        }
    )

    # Sample competitor price update data (different branch_codes to test item_code only join)
    competitor_prices = pl.DataFrame(
        {
            "competitor_id": ["COMP_001", "COMP_001"],
            "channel_type": ["OFFLINE", "OFFLINE"],
            "branch_code": ["CN04", "GN01"],  # Different branch codes from retail
            "item_code": ["7224000894", "7224000894"],
            "observed_at": ["2025-12-26 00:00:00", "2025-12-26 00:00:00"],
            "price_krw": ["4,660 ₩", "5,500 ₩"],
            "price_type": ["REGULAR", "REGULAR"],
            "promo_flag": [False, False],
            "source": ["TBD", "TBD"],
            "currency": ["KRW", "KRW"],
        }
    )

    print("=== Original Retail Data ===")
    print(
        retail.select(["id", "branch_code", "item_code", "competitor_price_krw_synth"])
    )

    print("\n=== Competitor Price Updates ===")
    print(competitor_prices.select(["branch_code", "item_code", "price_krw"]))

    result = run(
        retail_inventory=retail, competitor_prices=competitor_prices, options={}
    )

    print("\n=== Updated Retail Data ===")
    print(
        result["update_competitor_price"].select(
            ["id", "branch_code", "item_code", "competitor_price_krw_synth"]
        )
    )
