from typing import Optional, Any

import phonenumbers
from snowflake.snowpark.types import StringType
from snowflake.snowpark.functions import col, udf
from snowflake.snowpark import Session, DataFrame


def model(dbt: Any, session: Session) -> DataFrame:
    @udf(
        name="format_phone_number",
        return_type=StringType(),
        input_types=[StringType(), StringType()],
        replace=True,
    )
    def format_phone_number(number: str, region: Optional[str]) -> Optional[str]:
        try:
            phone_obj = phonenumbers.parse(number, region=region)
        except phonenumbers.NumberParseException:
            pass
        else:
            if phonenumbers.is_valid_number(phone_obj):
                return phonenumbers.format_number(
                    phone_obj,
                    num_format=phonenumbers.PhoneNumberFormat.E164,
                )
        return None

    fan_data_ref = dbt.ref("STG_EXTERNAL_EVENT_DBT")
    location_ref = dbt.ref("INT_FAN_LOCATION_DBT")

    sql = f"""
    SELECT DISTINCT FDR.FAN_ID, FDR.FAN_MOBILE, FDR.HAPPENED_AT, FL.COUNTRY_ISO2
    FROM {fan_data_ref.table_name} FDR
    LEFT JOIN {location_ref.table_name} FL
        ON FDR.FAN_ID = FL.FAN_ID
    WHERE FDR.FAN_MOBILE IS NOT NULL
    """

    return (
        session.sql(sql)
        .cache_result()
        .withColumn(
            "FAN_MOBILE_PHONE_ENRICHED",
            format_phone_number(
                col("FAN_MOBILE"),
                col("COUNTRY_ISO2"),
            ),
        )
        .filter(
            col("FAN_MOBILE_PHONE_ENRICHED").is_not_null(),
        )
        .select(
            col("FAN_ID"),
            col("HAPPENED_AT"),
            col("FAN_MOBILE_PHONE_ENRICHED").alias("FAN_MOBILE"),
        ).distinct()
    )