import polars as pl


def run(tb_sap_kna1, options, contexts):
    # 불필요한 칼럼 제거
    df = tb_sap_kna1.drop("mandt", "recordstamp")

    # 칼럼명 정리 및 나라 정보 mapping
    country_map = {
        "DE": "Germany",
        "FR": "France",
        "IT": "Italy",
        "ES": "Spain",
        "NL": "Netherlands",
        "US": "United States",
        "CA": "Canada",
        "MX": "Mexico",
        "KR": "South Korea",
        "JP": "Japan",
        "CN": "China",
        "SG": "Singapore",
    }

    anred_entity_type_mapping = {
        "Company": "organization/company",
        "Firma": "organization/company",
        "Mr.": "person",
        "Dr.": "person",
        "Frau": "person",
        "unknown": "unknown",
    }

    result = df.select(
        [
            pl.col("kunnr").alias("customer_id"),
            pl.col("name1").alias("customer_name"),
            pl.col("land1").alias("country_code"),
            pl.col("ort01").alias("city"),
            pl.col("regio").alias("region"),
            pl.col("anred").alias("from_address"),
        ]
    ).with_columns(
        [
            pl.when(pl.col("country_code").is_in(["DE", "FR", "IT", "ES", "NL"]))
            .then(pl.lit("europe"))
            .when(pl.col("country_code").is_in(["US", "CA", "MX"]))
            .then(pl.lit("north_america"))
            .when(pl.col("country_code").is_in(["KR", "JP", "CN", "SG"]))
            .then(pl.lit("asia"))
            .otherwise(pl.lit("other"))
            .alias("continent_name"),
            # country 이름 매핑
            pl.col("country_code").replace(country_map).alias("country_name"),
            # 고객 유형 매핑
            pl.col("from_address")
            .replace(anred_entity_type_mapping)
            .alias("customer_type"),
        ]
    )

    return {"tb_customer": result}
