# 检查特征工程 SDK 是否已安装（机器学习镜像已预装，会打印版本）
%pip list | grep wedata3-feature-engineering

# 未安装时取消下一行注释执行（要求内核 Python ≥ 3.10、MLflow ≥ 3.0.0），装完重启内核
# %pip install "wedata3-feature-engineering[mlflow3]==0.2.3"

import pandas as pd
from pyspark.sql.functions import current_timestamp
from wedata.feature_engineering.client import FeatureEngineeringClient
from wedata.feature_store.constants.engine_types import EngineTypes

# Studio（DLC）Notebook 内核已内置 spark 会话，可直接使用 spark 变量，无需手动创建 SparkSession

# 初始化特征工程客户端（鉴权信息由 Notebook 内核从环境变量自动注入，无需手动配置）
fe = FeatureEngineeringClient(spark)

# 读取上传的 CSV，并补充时间戳列
# 注意：请把下方路径替换为您在步骤 1.1 上传 CSV 的实际目录路径。
# 可在左侧文件管理器中右键该 CSV 文件、选择「复制路径」获取完整路径，
# 形如 /Workspace/Users/<你的账号>/customer_churn_class.csv
pdf = pd.read_csv("/Workspace/Users/<你的账号>/customer_churn_class.csv")
df = spark.createDataFrame(pdf).withColumn("event_timestamp", current_timestamp())

# 创建特征表：指定主键与时间戳键，engine_type 固定用 Iceberg，df 用于推断 schema 并写入首批数据
fe.create_table(
    name="customer_churn_features",          # （可修改）特征表名
    catalog_name="your_catalog",          # 目标 Catalog 名
    database_name="your_schema",           # 目标 Schema 名
    primary_keys=["customer_id"],            # 主键（支持复合主键）
    timestamp_key="event_timestamp",         # 时间戳键（单数）
    engine_type=EngineTypes.ICEBERG_ENGINE,
    df=df,                                    # 用 df 推断 schema 并写入首批数据
    description="客户流失特征表，用于 AutoML 训练。",
)

# 读回离线特征表，验证写入结果（表名为上面 create_table 的三段式全名，请替换为您自己的 catalog.schema）
fe.read_table(name="your_catalog.your_schema.customer_churn_features").show(5)

curl -X POST "https://<service-url>/predict" \\
  -H "Authorization: Bearer <token>" \\
  -H "Content-Type: application/json" \\
  -d '{"dataframe_records": [{
            "age": 27,
            "gender": "Male",
            "tenure_months": 70,
            "monthly_charges": 112.99,
            "total_charges": 7410.38,
            "contract_type": "One Year",
            "payment_method": "Credit Card",
            "num_support_tickets": 5,
            "avg_monthly_usage": 271.91
        }]}'

%pip install "wedata3-mlflow-tcdeploy-plugin"

from mlflow.deployments import get_deploy_client

# target 固定为 "tcdeploy"
client = get_deploy_client("tcdeploy")

# 控制台部署的服务，可通过 list_deployments() 获取服务组 ID（返回项的 name 即 service_group_id）
service_group_id = client.list_deployments()[0]["name"]

# predict 的 deployment_name 传服务组 ID（service_group_id，形如 msg-xxxxxx），不是服务名称
result = client.predict(
    deployment_name=service_group_id,
    inputs={"dataframe_records": [{
            "age": 27,
            "gender": "Male",
            "tenure_months": 70,
            "monthly_charges": 112.99,
            "total_charges": 7410.38,
            "contract_type": "One Year",
            "payment_method": "Credit Card",
            "num_support_tickets": 5,
            "avg_monthly_usage": 271.91
        }]},
   config={"auth_token": "Bearer <token>"},   # 服务开启鉴权时必填
)
print(result)