from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, trim

spark = SparkSession.builder \
    .appName("Clean-Full-Dataset-Fixed") \
    .master("spark://10.0.0.15:7077") \
    .getOrCreate()

spark.sparkContext.setLogLevel("WARN")

files = ["02-14-2018.csv", "02-15-2018.csv", "02-16-2018.csv",
         "02-20-2018.csv", "02-21-2018.csv", "02-22-2018.csv",
         "02-23-2018.csv", "02-28-2018.csv", "03-01-2018.csv",
         "03-02-2018.csv"]

extra_cols_in_0220 = ["Flow ID", "Src IP", "Src Port", "Dst IP"]

dfs = []
for f in files:
    path = "hdfs://10.0.0.15:9000/ids/raw_dataset/" + f
    d = spark.read.csv(path, header=True, inferSchema=True)
    # أسقطي الأعمدة الزائدة إن وُجدت في هذا الملف تحديداً
    cols_to_drop = [c for c in extra_cols_in_0220 if c in d.columns]
    if cols_to_drop:
        print(f + " -> dropping extra columns: " + str(cols_to_drop))
        d = d.drop(*cols_to_drop)
    dfs.append(d)

# دمج بمطابقة الأعمدة بالاسم (الأهم هنا)
df = dfs[0]
for d in dfs[1:]:
    df = df.unionByName(d, allowMissingColumns=False)

print("Raw combined rows: " + str(df.count()))

# استبعاد صفوف الـ header المكررة داخل الجسم (إن وجدت)
df = df.filter(col("Label") != "Label")

# تنظيف القيم اللانهائية/الفارغة
numeric_cols = [f.name for f in df.schema.fields
                if str(f.dataType) in ("DoubleType", "FloatType")]
for c in numeric_cols:
    df = df.filter(~col(c).isin(float("inf"), float("-inf")))
df = df.na.drop()

# تأكيد أن Label نص صالح فقط (استبعاد أي قيمة رقمية متسربة أخرى)
df = df.filter(~col("Label").rlike("^[0-9]"))

df = df.withColumn("Label_Binary",
        when(trim(col("Label")) == "Benign", 0).otherwise(1))

df.write.mode("overwrite").parquet("hdfs://10.0.0.15:9000/ids/dataset_full_clean.parquet")

print("Total clean rows: " + str(df.count()))
df.groupBy("Label").count().orderBy("count", ascending=False).show(20, truncate=False)

spark.stop()
