"""
Retrain Model — Human-in-the-loop retraining pipeline.

Joins analyst-reviewed alert verdicts (PostgreSQL) with the raw flow
features archived by ids_streaming.py (HDFS), builds a fresh balanced
training set, retrains the hierarchical two-stage model, and only
promotes the new model to production if it performs at least as well
as the current one on a held-out test split (champion-challenger).

Safety: runs in DRY_RUN mode by default (RETRAIN_DRY_RUN=true). In dry
run, the script computes everything and prints what WOULD happen, but
never overwrites the production model files. Set RETRAIN_DRY_RUN=false
to actually promote a winning challenger model.
"""

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, rand, when
from pyspark.ml.feature import VectorAssembler, StringIndexer
from pyspark.ml.classification import LogisticRegression, RandomForestClassifier
from pyspark.ml.evaluation import MulticlassClassificationEvaluator
import psycopg2
import json
import os
from datetime import datetime

# ── Config ──
HDFS_ARCHIVE = "hdfs://10.0.0.15:9000/ids/model_results"
MODEL_BASE = "/models"
REGISTRY_PATH = "/models/model_registry.json"

DB_CONFIG = dict(
    host="10.0.0.18", port=5432,
    dbname="ids_db", user="ids_user", password=os.getenv("DB_PASSWORD", "")
)

DRY_RUN = os.getenv("RETRAIN_DRY_RUN", "true").lower() == "true"

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"
]

spark = SparkSession.builder \
    .appName("IDS-Retrain") \
    .master("spark://10.0.0.15:7077") \
    .config("spark.driver.host", "10.0.0.15") \
    .config("spark.driver.bindAddress", "0.0.0.0") \
    .getOrCreate()

spark.sparkContext.setLogLevel("WARN")


def log(msg):
    print(f"[RETRAIN] {msg}", flush=True)


# ── 1. Pull analyst-reviewed alerts from PostgreSQL ──
def fetch_reviewed_alerts():
    conn = psycopg2.connect(**DB_CONFIG)
    cur = conn.cursor()
    cur.execute("""
        SELECT src_ip, attack_type, timestamp, verdict
        FROM alerts
        WHERE verdict IN ('true_positive', 'false_positive')
    """)
    rows = cur.fetchall()
    cur.close()
    conn.close()
    log(f"Fetched {len(rows)} analyst-reviewed alerts")
    return rows


# ── 2. Join reviewed alerts with HDFS raw features ──
def build_reviewed_dataset(spark, reviewed_rows):
    if not reviewed_rows:
        return None

    reviewed_df = spark.createDataFrame(
        reviewed_rows, ["src_ip", "attack_type", "timestamp", "verdict"]
    )

    archive = spark.read.parquet(f"{HDFS_ARCHIVE}/*")

    joined = archive.join(
        reviewed_df,
        on=["src_ip", "timestamp"],
        how="inner"
    )

    # Ground truth: confirmed attacks keep their type, false positives become Benign
    labeled = joined.withColumn(
        "ground_truth",
        when(col("verdict") == "true_positive", col("attack_type"))
        .otherwise("Benign")
    )

    return labeled.select(*feature_cols, "ground_truth")


# ── 3. Add bulk Benign samples from HDFS archive for balance ──
def add_benign_pool(spark, reviewed_dataset):
    archive = spark.read.parquet(f"{HDFS_ARCHIVE}/*")
    benign_pool = archive.filter(col("predicted_label") == "Benign") \
                          .select(*feature_cols, col("predicted_label").alias("ground_truth"))

    reviewed_count = reviewed_dataset.count() if reviewed_dataset else 0
    benign_needed = max(reviewed_count, 500)

    benign_sample = benign_pool.orderBy(rand()).limit(benign_needed)

    if reviewed_dataset is not None:
        combined = reviewed_dataset.unionByName(benign_sample)
    else:
        combined = benign_sample

    return combined


# ── 4. Model registry (tracks current production model's accuracy) ──
def load_registry():
    if os.path.exists(REGISTRY_PATH):
        with open(REGISTRY_PATH) as f:
            return json.load(f)
    return {"version": 0, "accuracy": 0.0, "trained_at": None}


def save_registry(data):
    with open(REGISTRY_PATH, "w") as f:
        json.dump(data, f, indent=2)


# ── 5. Train + Evaluate ──
def train_and_evaluate(dataset):
    dataset = dataset.dropna()
    assembler = VectorAssembler(inputCols=feature_cols, outputCol="features", handleInvalid="skip")
    indexer = StringIndexer(inputCol="ground_truth", outputCol="label").fit(dataset)

    prepared = indexer.transform(assembler.transform(dataset))
    train_df, test_df = prepared.randomSplit([0.8, 0.2], seed=42)

    log(f"Training on {train_df.count()} rows, testing on {test_df.count()} rows")

    # Stage 1: binary (Benign vs Attack)
    train_bin = train_df.withColumn("binary_label", when(col("ground_truth") == "Benign", 0.0).otherwise(1.0))
    test_bin = test_df.withColumn("binary_label", when(col("ground_truth") == "Benign", 0.0).otherwise(1.0))

    lr = LogisticRegression(featuresCol="features", labelCol="binary_label", maxIter=50)
    stage1_model = lr.fit(train_bin)

    # Stage 2: multiclass among attack types only
    train_attacks = train_df.filter(col("ground_truth") != "Benign")
    test_attacks = test_df.filter(col("ground_truth") != "Benign")

    stage2_model = rf_fit_if_possible(train_attacks)

    # ── Evaluate full pipeline on the held-out test set ──
    bin_preds = stage1_model.transform(test_bin)
    evaluator_bin = MulticlassClassificationEvaluator(
        labelCol="binary_label", predictionCol="prediction", metricName="accuracy"
    )
    stage1_acc = evaluator_bin.evaluate(bin_preds)

    overall_acc = stage1_acc
    if stage2_model is not None and test_attacks.count() > 0:
        multi_preds = stage2_model.transform(test_attacks)
        evaluator_multi = MulticlassClassificationEvaluator(
            labelCol="label", predictionCol="prediction", metricName="accuracy"
        )
        stage2_acc = evaluator_multi.evaluate(multi_preds)
        overall_acc = (stage1_acc + stage2_acc) / 2
        log(f"Stage1 accuracy: {stage1_acc:.4f} | Stage2 accuracy: {stage2_acc:.4f}")
    else:
        log(f"Stage1 accuracy: {stage1_acc:.4f} | Stage2 skipped (no attack samples)")

    return stage1_model, stage2_model, indexer, overall_acc


def rf_fit_if_possible(train_attacks):
    if train_attacks.count() == 0:
        return None
    rf = RandomForestClassifier(featuresCol="features", labelCol="label", numTrees=100)
    return rf.fit(train_attacks)


# ── Main ──
def main():
    log(f"Starting retraining pipeline... (DRY_RUN={DRY_RUN})")

    reviewed_rows = fetch_reviewed_alerts()
    reviewed_dataset = build_reviewed_dataset(spark, reviewed_rows) if reviewed_rows else None
    full_dataset = add_benign_pool(spark, reviewed_dataset)

    total_rows = full_dataset.count()
    log(f"Total training rows after balancing: {total_rows}")

    if total_rows < 50:
        log("Not enough data to retrain yet. Aborting.")
        print("RETRAIN_RESULT: insufficient_data")
        spark.stop()
        return

    stage1_model, stage2_model, indexer, new_accuracy = train_and_evaluate(full_dataset)

    registry = load_registry()
    old_accuracy = registry.get("accuracy", 0.0)

    log(f"New model accuracy: {new_accuracy:.4f} | Current production accuracy: {old_accuracy:.4f}")

    if new_accuracy >= old_accuracy:
        if DRY_RUN:
            log("New model WOULD be promoted (better or equal) — DRY_RUN=true, no files changed.")
            print(f"RETRAIN_RESULT: would_promote new_accuracy={new_accuracy:.4f} old_accuracy={old_accuracy:.4f}")
        else:
            log("New model is better or equal — promoting to production.")
            stage1_model.write().overwrite().save(f"{MODEL_BASE}/hierarchical_stage1_model")
            if stage2_model is not None:
                stage2_model.write().overwrite().save(f"{MODEL_BASE}/hierarchical_stage2_model")
            indexer.write().overwrite().save(f"{MODEL_BASE}/hierarchical_indexer_model")

            registry = {
                "version": registry.get("version", 0) + 1,
                "accuracy": new_accuracy,
                "trained_at": datetime.now().isoformat(),
                "training_rows": total_rows
            }
            save_registry(registry)
            print(f"RETRAIN_RESULT: promoted accuracy={new_accuracy:.4f}")
    else:
        log("New model performs worse — keeping current production model (no action needed).")
        print(f"RETRAIN_RESULT: rejected new_accuracy={new_accuracy:.4f} old_accuracy={old_accuracy:.4f}")

    spark.stop()


if __name__ == "__main__":
    main()
