import csv
import os
import time
from datetime import datetime
from hdfs import InsecureClient

SOURCE_DIR      = "/source_attacks"
OUTPUT_DIR      = "/output"
ROWS_PER_DRIP   = 500
DRIP_INTERVAL   = 3
TIMESTAMP_COL_INDEX = 5
LOOP_WHEN_DONE  = True

HDFS_URL        = "http://10.0.0.15:9870"
HDFS_USER       = "root"
HDFS_ARCHIVE_BASE = "/ids/archive"

os.makedirs(OUTPUT_DIR, exist_ok=True)
hdfs_client = InsecureClient(HDFS_URL, user=HDFS_USER)


def get_source_files():
    return sorted([
        f for f in os.listdir(SOURCE_DIR)
        if f.endswith(".csv")
    ])


def is_file_stable(path, wait=2):
    """Make sure the file isn't still being written (size unchanged after `wait` seconds)."""
    size1 = os.path.getsize(path)
    time.sleep(wait)
    size2 = os.path.getsize(path)
    return size1 == size2 and size1 > 0


def drip_file(header, rows, batch_num):
    now_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    out_rows = []
    for row in rows:
        row = row.copy()
        if len(row) > TIMESTAMP_COL_INDEX:
            row[TIMESTAMP_COL_INDEX] = now_str
        out_rows.append(row)

    fname = f"drip_{datetime.now().strftime('%Y%m%d_%H%M%S_%f')}.csv"
    out_path = os.path.join(OUTPUT_DIR, fname)

    with open(out_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(header)
        writer.writerows(out_rows)

    print(f"[{now_str}] Dripped {len(out_rows)} rows -> {fname}", flush=True)


def archive_to_hdfs(local_path, filename):
    today = datetime.now()
    hdfs_dir = f"{HDFS_ARCHIVE_BASE}/{today.year}/{today.month:02d}/{today.day:02d}"
    hdfs_path = f"{hdfs_dir}/{filename}"
    try:
        hdfs_client.makedirs(hdfs_dir)
        hdfs_client.upload(hdfs_path, local_path, overwrite=True)
        print(f"[HDFS] Archived -> {hdfs_path}", flush=True)
        return True
    except Exception as e:
        print(f"[HDFS ERROR] Failed to archive {filename}: {e}", flush=True)
        return False


def process_source_file(fname):
    path = os.path.join(SOURCE_DIR, fname)
    marker_path = path + ".dripped"   # علامة: تم التفريغ فعلاً

    if not is_file_stable(path):
        print(f"[SKIP] {fname} is empty or still being written, will retry next cycle", flush=True)
        return

    # لو أصلاً تفرّغ بالكامل قبل هيك، بس الأرشفة فشلت → لا تعيدي الـ drip، حاولي الأرشفة بس
    if os.path.exists(marker_path):
        print(f"[RETRY ARCHIVE] {fname} already dripped, retrying HDFS archive only", flush=True)
        if archive_to_hdfs(path, fname):
            os.remove(path)
            os.remove(marker_path)
            print(f"[CLEANUP] Removed {fname} from {SOURCE_DIR}", flush=True)
        return

    try:
        with open(path, "r") as f:
            reader = csv.reader(f)
            header = next(reader)
            rows = list(reader)
    except StopIteration:
        print(f"[SKIP] {fname} has no data rows, skipping", flush=True)
        return
    except Exception as e:
        print(f"[ERROR] Failed to read {fname}: {e}", flush=True)
        return

    if not rows:
        print(f"[SKIP] {fname} has header only, no rows to drip", flush=True)
        return

    print(f"Loaded {len(rows)} rows from {fname}", flush=True)

    for i in range(0, len(rows), ROWS_PER_DRIP):
        chunk = rows[i:i + ROWS_PER_DRIP]
        drip_file(header, chunk, i // ROWS_PER_DRIP)
        time.sleep(DRIP_INTERVAL)

    
    open(marker_path, "w").close()

    if archive_to_hdfs(path, fname):
        try:
            os.remove(path)
            os.remove(marker_path)
            print(f"[CLEANUP] Removed {fname} from {SOURCE_DIR}", flush=True)
        except Exception as e:
            print(f"[ERROR] Could not remove {fname}: {e}", flush=True)

def main():
    print("Traffic Dripper (containerized) started.", flush=True)
    while True:
        source_files = get_source_files()
        if not source_files:
            print(f"No CSV files found in {SOURCE_DIR}. Waiting...", flush=True)
            time.sleep(30)
            continue

        for fname in source_files:
            try:
                process_source_file(fname)
            except Exception as e:
                print(f"[FATAL for {fname}] {e} — skipping this file, continuing loop", flush=True)
                continue

        if not LOOP_WHEN_DONE:
            print("All source files processed. Exiting.", flush=True)
            break

        print("Finished one pass. Looping...", flush=True)
        time.sleep(5)


if __name__ == "__main__":
    main()
