from pyspark.sql import SparkSession
from pyspark.sql.functions import from_json, col, lit, udf, count as spark_count, max as spark_max
from pyspark.sql.types import StructType, StructField, StringType, DoubleType
from pyspark.ml.classification import LogisticRegressionModel, RandomForestClassificationModel
from pyspark.ml.feature import VectorAssembler, StringIndexerModel
import json
import requests
from datetime import datetime
from pyspark.sql.functions import from_json, col, lit, udf, count as spark_count, max as spark_max, first


spark = SparkSession.builder \
    .appName("Smart-IDS-Streaming") \
    .master("local[2]") \
    .getOrCreate()

spark.sparkContext.setLogLevel("WARN")
print("Spark Session started")

print("Loading AI models...")
model_stage1  = LogisticRegressionModel.load("/models/hierarchical_stage1_model")
model_stage2  = RandomForestClassificationModel.load("/models/hierarchical_stage2_model")
indexer_model = StringIndexerModel.load("/models/hierarchical_indexer_model")
print("Models loaded")

feature_cols = [
    "Dst Port", "Protocol", "Flow Duration",
    "Tot Fwd Pkts", "Tot Bwd Pkts",
    "TotLen Fwd Pkts", "TotLen Bwd Pkts",
    "Fwd Pkt Len Max", "Fwd Pkt Len Min",
    "Bwd Pkt Len Max", "Bwd Pkt Len Min",
    "Flow Pkts/s", "Flow Byts/s",
    "Fwd Header Len", "Bwd Header Len",
    "Fwd Pkts/s", "Bwd Pkts/s",
    "Pkt Len Min", "Pkt Len Max",
    "SYN Flag Cnt", "RST Flag Cnt",
    "PSH Flag Cnt", "ACK Flag Cnt"
]

schema_fields = [
    StructField("timestamp", StringType(), True),
    StructField("src_ip", StringType(), True),
    StructField("label", StringType(), True),
]
for f in feature_cols:
    schema_fields.append(StructField(f, DoubleType(), True))

log_schema = StructType(schema_fields)

assembler = VectorAssembler(
    inputCols=feature_cols,
    outputCol="features",
    handleInvalid="skip"
)

# Maps a Stage 2 numeric class index back to its human-readable attack
# name, using the labels learned by the StringIndexer during training.
model_labels = indexer_model.labels

def index_to_label(idx):
    idx = int(idx)
    return model_labels[idx] if idx < len(model_labels) else "Unknown"

index_to_label_udf = udf(index_to_label, StringType())

# HDFS path where every processed batch is archived as Parquet,
# including the raw features used for the prediction. This archive is
# later joined with the analyst verdicts stored in PostgreSQL
# (matched on src_ip + timestamp) to build a labeled retraining set.
HDFS_RESULTS_PATH = "hdfs://namenode:9000/ids/model_results"

# ── NEW: Blocked IPs Filter ──
# Before analyzing any batch, Spark asks FastAPI which IPs are
# currently blocked, and drops their logs BEFORE running them through
# the model. This avoids wasting inference cycles on traffic that's
# already handled, and avoids spamming duplicate alerts for a source
# the admin already blocked.
BLOCKED_IPS_ENDPOINT = "http://10.0.0.124:8000/api/blocked-ips"

def get_blocked_ips():
    try:
        resp = requests.get(BLOCKED_IPS_ENDPOINT, timeout=3)
        if resp.status_code == 200:
            data = resp.json()
            # Adjust parsing here if the FastAPI response shape differs,
            # e.g. {"blocked_ips": [...]} vs a plain list [...]
            if isinstance(data, dict):
                return set(data.get("blocked_ips", []))
            return set(data)
    except Exception as ex:
        print("Could not fetch blocked IPs, proceeding without filter: " + str(ex))
    return set()

print("Connecting to Kafka...")
raw_stream = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "kafka:9092") \
    .option("subscribe", "network-logs") \
    .option("startingOffsets", "earliest") \
    .load()

parsed = raw_stream.select(
    from_json(col("value").cast("string"), log_schema).alias("data")
).select("data.*")

def process_batch(batch_df, batch_id):
    if batch_df.isEmpty():
        return
    count = batch_df.count()
    print("=== Batch " + str(batch_id) + " --- " + str(count) + " logs ===")

    # ── NEW: filter out already-blocked source IPs before analysis ──
    blocked_ips = get_blocked_ips()
    if blocked_ips:
        before_count = count
        batch_df = batch_df.filter(~col("src_ip").isin(list(blocked_ips)))
        count = batch_df.count()
        print("Filtered " + str(before_count - count) + " logs from " +
              str(len(blocked_ips)) + " blocked IP(s). Remaining: " + str(count))
        if batch_df.isEmpty():
            print("All logs in this batch were from blocked IPs — skipping batch.")
            return

    try:
        featured = assembler.transform(batch_df)
        s1 = model_stage1.transform(featured) \
                         .withColumnRenamed("prediction", "stage1_prediction")
        normal_df = s1.filter(col("stage1_prediction") == 0.0)
        attack_df = s1.filter(col("stage1_prediction") == 1.0)
        normal_count = normal_df.count()
        attack_count = attack_df.count()
        print("Normal: " + str(normal_count) + " | Attacks: " + str(attack_count))

        # Build the archive rows for normal traffic (predicted_label = Benign)
        normal_out = normal_df.select(
            *feature_cols, "timestamp", "src_ip"
        ).withColumn("predicted_label", lit("Benign"))

        attack_out = None
        if attack_count > 0:
            s2 = model_stage2.transform(attack_df.drop("stage1_prediction")) \
                             .withColumnRenamed("prediction", "stage2_prediction")

            # Build the archive rows for attack traffic, with the
            # human-readable predicted attack type. This keeps every
            # individual flow in the HDFS archive (needed for retraining
            # granularity), even though alerts sent to FastAPI are
            # aggregated below.
            attack_out = s2.withColumn(
                "predicted_label", index_to_label_udf(col("stage2_prediction"))
            ).select(*feature_cols, "timestamp", "src_ip", "predicted_label")

            # ── Aggregate by (src_ip, attack_type) within this micro-batch ──
            # so e.g. 100 flows from the same source in the same 15s window
            # become ONE alert with an occurrence_count, not 100 separate
            # alerts. This is what fixes the alert flood on the dashboard.
            aggregated = s2.withColumn(
             "attack_type_resolved", index_to_label_udf(col("stage2_prediction"))
        ).groupBy("src_ip", "attack_type_resolved").agg(
            spark_count("*").alias("occurrence_count"),
            spark_max("timestamp").alias("timestamp"),
            *[first(c).alias(c) for c in feature_cols]
        )

            attacks = aggregated.collect()
            for row in attacks:
                attack_type = row["attack_type_resolved"]
                src_ip = row["src_ip"]
                occurrence_count = row["occurrence_count"]

                sample_features = {c: row[c] for c in feature_cols}

                alert = {
                    "attack_type": attack_type,
                    "src_ip": src_ip,
                    "timestamp": row["timestamp"],
                    "severity": "HIGH" if "DDoS" in attack_type else "MEDIUM",
                    "occurrence_count": occurrence_count,
                    "sample_features": sample_features
                }

                print("ALERT: " + json.dumps(alert))
                try:
                    requests.post(
                        "http://10.0.0.124:8000/api/alerts",
                        json=alert,
                        timeout=3
                    )
                except Exception as ex:
                    print("FastAPI error: " + str(ex))

        # Combine normal + attack rows into one archive batch.
        # NOTE: this still archives every individual flow (not aggregated)
        # to HDFS, since retraining needs per-flow feature granularity.
        if attack_out is not None:
            archive_df = normal_out.unionByName(attack_out)
        else:
            archive_df = normal_out

        # Write this batch to HDFS as Parquet -- one file per micro-batch
        # (~every 15s), partitioned by ingestion date to keep the
        # directory organized as the archive grows.
        ingest_date = datetime.now().strftime("%Y-%m-%d")
        out_path = HDFS_RESULTS_PATH + "/date=" + ingest_date
        archive_df.write.mode("append").parquet(out_path)
        print("Archived " + str(archive_df.count()) + " rows to HDFS: " + out_path)

    except Exception as e:
        print("Batch error: " + str(e))

print("Starting Spark Streaming...")
query = parsed.writeStream \
    .foreachBatch(process_batch) \
    .trigger(processingTime="15 seconds") \
    .option("checkpointLocation", "/tmp/spark_checkpoint") \
    .start()

query.awaitTermination()