from flask import Flask, jsonify, send_file, request
from flask_cors import CORS
import threading
import multiprocessing as mp
import gc
import traceback
import json
import time
import os
import platform
import psutil
import traceback
try:
    import winreg
except ImportError:
    winreg = None
from datetime import datetime
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import A4
from reportlab.lib.units import cm

from isolation_forest_cicids import run_isolation_forest
from autoencoder import run_autoencoder
from isolation_forest_custom import run_isolation_forest_custom
from autoencoder_custom import run_autoencoder_custom
import signal
import sys
import shutil
from flask import send_from_directory

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

TEMP_DIR = os.path.join(BASE_DIR, "temp")
REPORT_DIR = os.path.join(BASE_DIR, "reports")
PROGRESS_PATH = os.path.join(BASE_DIR, "progress.json")
FRONTEND_DIR = os.path.join(BASE_DIR, "..", "Frontend")
HTML_DIR = os.path.join(FRONTEND_DIR, "HTML Files")
JS_DIR = os.path.join(FRONTEND_DIR, "JS")
IMG_DIR = os.path.join(FRONTEND_DIR, "img")

os.makedirs(TEMP_DIR, exist_ok=True)
os.makedirs(REPORT_DIR, exist_ok=True)

app = Flask(__name__)
CORS(app)

def force_memory_cleanup():
    gc.collect()
    try:
        import ctypes
        ctypes.CDLL("libc.so.6").malloc_trim(0)
    except Exception:
        pass

@app.route("/", methods=["GET"])
def serve_index():
    return send_from_directory(HTML_DIR, "index.html")

@app.route("/upload", methods=["GET"])
def serve_upload():
    return send_from_directory(HTML_DIR, "upload.html")

@app.route("/default", methods=["GET"])
def serve_default():
    return send_from_directory(HTML_DIR, "default.html")

@app.route("/JS/<path:filename>", methods=["GET"])
def serve_js(filename):
    return send_from_directory(JS_DIR, filename)

@app.route("/img/<path:filename>", methods=["GET"])
def serve_img(filename):
    return send_from_directory(IMG_DIR, filename)

model_result_iforest = None
model_result_autoencoder = None
model_result_iforest_custom = None
model_result_autoencoder_custom = None

state_lock = threading.Lock()
is_running = False
current_job = None
worker_process = None

RESULT_IFOREST = os.path.join(TEMP_DIR, "result_iforest.json")
RESULT_AUTOENCODER = os.path.join(TEMP_DIR, "result_autoencoder.json")
RESULT_IFOREST_CUSTOM = os.path.join(TEMP_DIR, "result_iforest_custom.json")
RESULT_AUTOENCODER_CUSTOM = os.path.join(TEMP_DIR, "result_autoencoder_custom.json")

def handle_exit(signum, frame):
    cleanup()
    sys.exit(0)

signal.signal(signal.SIGINT, handle_exit)
signal.signal(signal.SIGTERM, handle_exit)


def clear_directory(path):
    if not os.path.exists(path):
        return

    for item in os.listdir(path):
        item_path = os.path.join(path, item)
        try:
            if os.path.isfile(item_path) or os.path.islink(item_path):
                os.unlink(item_path)
            elif os.path.isdir(item_path):
                shutil.rmtree(item_path)
        except Exception as e:
            print(f"Failed to delete {item_path}: {e}")

def cleanup():

    clear_directory(TEMP_DIR)
    clear_directory(REPORT_DIR)


def update_progress(value):
    with open(PROGRESS_PATH, "w") as f:
        json.dump({"progress": value}, f)

def read_json_file(path, default=None):
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return default


def write_json_file(path, data):
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)


def get_result_path(job_name):
    mapping = {
        "iforest": RESULT_IFOREST,
        "autoencoder": RESULT_AUTOENCODER,
        "iforest_custom": RESULT_IFOREST_CUSTOM,
        "autoencoder_custom": RESULT_AUTOENCODER_CUSTOM,
    }
    return mapping[job_name]


def get_job_runner(job_name):
    if job_name == "iforest":
        return lambda: run_isolation_forest(plot=False, table=False)
    elif job_name == "autoencoder":
        return lambda: run_autoencoder(plot=False, table=False)
    elif job_name == "iforest_custom":
        return run_isolation_forest_custom
    elif job_name == "autoencoder_custom":
        return run_autoencoder_custom
    else:
        raise ValueError(f"Unknown job name: {job_name}")


def worker_entry(job_name):
    result_path = get_result_path(job_name)
    start = time.time()

    try:
        update_progress(1)

        runner = get_job_runner(job_name)
        result = runner()

        end = time.time()
        if isinstance(result, dict):
            result["runtime"] = round(end - start, 2)

        write_json_file(result_path, {
            "ready": True,
            "data": result
        })

        update_progress(100)

    except Exception as e:
        write_json_file(result_path, {
            "ready": True,
            "data": {
                "error": True,
                "message": str(e),
                "details": traceback.format_exc()
            }
        })
        update_progress(0)

def refresh_worker_state():
    global is_running, worker_process, current_job

    with state_lock:
        proc = worker_process
        running = is_running

    if running and proc is not None and not proc.is_alive():
        try:
            proc.join(timeout=0.2)
        except Exception:
            pass

        try:
            proc.close()
        except Exception:
            pass

        with state_lock:
            is_running = False
            worker_process = None
            current_job = None

        force_memory_cleanup()



def start_background_job(job_name):
    global is_running, current_job, worker_process
    result_path = get_result_path(job_name)

    refresh_worker_state()

    with state_lock:
        if is_running:
            return False

        is_running = True
        current_job = job_name

    if os.path.exists(result_path):
        try:
            os.remove(result_path)
        except Exception:
            pass

    update_progress(0)

    proc = mp.Process(target=worker_entry, args=(job_name,))
    proc.start()

    with state_lock:
        worker_process = proc

    return True


def read_progress():
    try:
        with open(PROGRESS_PATH, "r") as f:
            return json.load(f).get("progress", 0)
    except Exception:
        return 0

update_progress(0)


@app.route("/start-iforest", methods=["POST"])
def start_iforest():
    started = start_background_job("iforest")
    if not started:
        return jsonify({
            "status": "already_running",
            "message": "Another algorithm is already running"
        }), 409
    return jsonify({"status": "started"})


@app.route("/start-iforest-custom", methods=["POST"])
def start_iforest_custom():
    started = start_background_job("iforest_custom")
    if not started:
        return jsonify({
            "status": "already_running",
            "message": "Another algorithm is already running"
        }), 409
    return jsonify({"status": "started"})

@app.route("/start-autoencoder-custom", methods=["POST"])
def start_autoencoder_custom():
    started = start_background_job("autoencoder_custom")
    if not started:
        return jsonify({
            "status": "already_running",
            "message": "Another algorithm is already running"
        }), 409
    return jsonify({"status": "started"})


@app.route("/start-autoencoder", methods=["POST"])
def start_autoencoder():
    started = start_background_job("autoencoder")
    if not started:
        return jsonify({
            "status": "already_running",
            "message": "Another algorithm is already running"
        }), 409
    return jsonify({"status": "started"})


@app.route("/get-progress", methods=["GET"])
def get_progress():
    refresh_worker_state()
    return jsonify(read_json_file(PROGRESS_PATH, {"progress": 0}))


@app.route("/get-status", methods=["GET"])
def get_status():
    refresh_worker_state()

    with state_lock:
        running = is_running
        job = current_job

    progress_data = read_json_file(PROGRESS_PATH, {"progress": 0})

    return jsonify({
        "running": running,
        "current_job": job,
        "progress": progress_data.get("progress", 0)
    })


@app.route("/get-isolation-forest-result", methods=["GET"])
def get_iforest_result():
    refresh_worker_state()
    data = read_json_file(RESULT_IFOREST)
    if not data:
        return jsonify({"ready": False})
    return jsonify(data)

@app.route("/get-isolation-forest-result-custom", methods=["GET"])
def get_iforest_result_custom():
    refresh_worker_state()
    data = read_json_file(RESULT_IFOREST_CUSTOM)
    if not data:
        return jsonify({"ready": False})
    return jsonify(data)


@app.route("/get-autoencoder-result", methods=["GET"])
def get_auto_result():
    refresh_worker_state()
    data = read_json_file(RESULT_AUTOENCODER)
    if not data:
        return jsonify({"ready": False})
    return jsonify(data)

@app.route("/get-autoencoder-result-custom", methods=["GET"])
def get_autoencoder_result_custom():
    refresh_worker_state()
    data = read_json_file(RESULT_AUTOENCODER_CUSTOM)
    if not data:
        return jsonify({"ready": False})
    return jsonify(data)



def generate_pdf_report(results, algorithm, params):
    filename = os.path.join(REPORT_DIR, "report.pdf")
    c = canvas.Canvas(filename, pagesize=A4)

    w, h = A4
    y = h - 2*cm


    c.setFont("Helvetica-Bold", 24)
    c.drawCentredString(w/2, y, f"{algorithm} Report")
    y -= 1.4*cm

    c.setFont("Helvetica", 11)
    c.drawCentredString(
        w/2, y,
        f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
    )
    y -= 1.6*cm


    c.setFont("Helvetica-Bold", 15)
    c.drawString(2*cm, y, "Algorithm Parameters:")
    y -= 0.9*cm

    c.setFont("Helvetica", 11)
    t = c.beginText(2.2*cm, y)
    t.setLeading(13)
    for k, v in params.items():
        t.textLine(f"{k}: {v}")
    c.drawText(t)
    y = t.getY() - 1.1*cm


    c.setFont("Helvetica-Bold", 15)
    c.drawString(2*cm, y, "Execution Summary:")
    y -= 0.9*cm

    c.setFont("Helvetica", 11)
    t = c.beginText(2.2*cm, y)
    t.setLeading(13)
    t.textLine(f"Runtime: {results['runtime']} seconds")
    t.textLine(f"Normal Traffic: {int(results['normal_count'])}")
    t.textLine(f"Detected Attacks: {int(results['attack_count'])}")
    t.textLine(f"Accuracy: {results['accuracy']:.4f}")
    c.drawText(t)
    y = t.getY() - 1.3*cm


    sys = get_system_info()

    c.setFont("Helvetica-Bold", 15)
    c.drawString(2 * cm, y, "System Information")
    y -= 0.9 * cm

    c.setFont("Helvetica", 11)
    info = c.beginText(2.2 * cm, y)
    info.setLeading(13)
    info.textLine(f"CPU: {sys['cpu']}")
    info.textLine(f"Physical Cores: {sys['cpu_physical_cores']}")
    info.textLine(f"Logical Cores: {sys['cpu_total_cores']}")
    info.textLine(f"GPU: {sys['gpu_name']}")
    info.textLine(f"OS: {sys['os']}")
    info.textLine(f"Python Version: {sys['python_version']}")
    info.textLine(f"System RAM: {sys['ram_total_gb']} GB")
    c.drawText(info)
    y = info.getY() - 1.3 * cm


    c.setFont("Helvetica-Bold", 15)
    c.drawString(2 * cm, y, "Memory Usage:")
    y -= 0.9 * cm

    c.setFont("Helvetica", 11)
    t = c.beginText(2.2 * cm, y)
    t.setLeading(13)
    t.textLine(f"RAM Before: {results['ram_before']} MB")
    t.textLine(f"Peak RAM: {results['ram_peak']} MB")
    t.textLine(f"RAM After: {results['ram_after']} MB")
    t.textLine(f"Total Increase: {results['ram_increase']} MB")
    c.drawText(t)
    y = t.getY() - 1.1 * cm


    c.setFont("Helvetica-Bold", 15)
    c.drawString(2*cm, y, "Top 5 Most Significant Anomalies:")
    y -= 1*cm

    c.setFont("Helvetica", 11)
    for anomaly in results["top_anomalies"]:
        t = c.beginText(2.4*cm, y)
        t.setLeading(13)
        for k, v in anomaly.items():
            t.textLine(f"{k}: {v}")
        c.drawText(t)
        y = t.getY() - 0.9*cm

        if y < 3*cm:
            c.showPage()
            y = h - 2*cm
            c.setFont("Helvetica", 11)

    c.save()
    return filename

@app.route("/download-pdf", methods=["POST"])
def download_pdf():
    data = request.get_json()
    algorithm = data["algorithm"]
    results = data["results"]

    raw_algorithm = data["algorithm"]
    algorithm_name = raw_algorithm.lower()

    if "isolation forest" in algorithm_name:
        algo_title = "Isolation Forest"
        params = {
            "n_estimators": 600,
            "max_samples": 0.3,
            "contamination": results.get("contamination"),
            "max_features": 0.7
        }

    elif "autoencoder" in algorithm_name:
        algo_title = "Autoencoder"
        params = {
            "epochs": 20,
            "batch_size": 64,
            "latent_dimension": 16,
            "validation_split": 0.1
        }

    else:
        params = {}

    pdf_path = generate_pdf_report(results, algo_title, params)
    return send_file(pdf_path, as_attachment=True)

@app.route("/get-system-info", methods=["GET"])
def system_info_endpoint():
    return jsonify(get_system_info())

def get_cpu_name():
    try:
        key = winreg.OpenKey(
            winreg.HKEY_LOCAL_MACHINE,
            r"HARDWARE\DESCRIPTION\System\CentralProcessor\0"
        )
        cpu_name, _ = winreg.QueryValueEx(key, "ProcessorNameString")
        winreg.CloseKey(key)
        return cpu_name.strip()
    except:
        return platform.processor() or "Unknown CPU"

def get_system_info():
    return {
        "cpu": get_cpu_name(),
        "cpu_physical_cores": psutil.cpu_count(logical=False),
        "cpu_total_cores": psutil.cpu_count(logical=True),
        "gpu_available": False,
        "gpu_name": "None",
        "os": platform.platform(),
        "python_version": platform.python_version(),
        "ram_total_gb": round(psutil.virtual_memory().total / (1024**3), 2),
        "ram_used_mb": round(psutil.virtual_memory().used / (1024**2), 2),
    }

@app.route("/download-json", methods=["POST"])
def download_json():
    data = request.get_json()

    algorithm = data["algorithm"]
    results = data["results"]
    runtime = data["runtime"]

    raw_algorithm = data["algorithm"]
    algorithm_name = raw_algorithm.lower()

    if "isolation forest" in algorithm_name:
        algo_title = "Isolation Forest"
        params = {
            "n_estimators": 600,
            "max_samples": 0.3,
            "contamination": results.get("contamination"),
            "max_features": 0.7
        }

    elif "autoencoder" in algorithm_name:
        algo_title = "Autoencoder"
        params = {
            "epochs": 20,
            "batch_size": 64,
            "latent_dimension": 16,
            "validation_split": 0.1
        }

    system_info = get_system_info()

    ram_info = {
        "ram_before": results["ram_before"],
        "ram_peak": results["ram_peak"],
        "ram_after": results["ram_after"],
        "ram_increase": results["ram_increase"]
    }

    anomalies = results.get("top_anomalies", [])

    json_report = {
        "algorithm": algorithm,
        "generated": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "algorithm_parameters": params,

        "execution_summary": {
            "runtime_seconds": runtime,
            "normal_traffic": results["normal_count"],
            "detected_attacks": results["attack_count"],
            "accuracy": results["accuracy"]
        },

        "system_information": system_info,

        "memory_usage": ram_info,

        "top_5_anomalies": anomalies
    }

    json_path = os.path.join(REPORT_DIR, "report.json")
    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(json_report, f, indent=4)

    return send_file(json_path, as_attachment=True)


@app.route("/upload-dataset", methods=["POST"])
def upload_dataset():
    refresh_worker_state()

    with state_lock:
        if is_running:
            return jsonify({
                "error": "Cannot upload a new dataset while an algorithm is running"
            }), 409

    if "file" not in request.files:
        return jsonify({"error": "No file provided"}), 400

    file = request.files["file"]

    if file.filename == "":
        return jsonify({"error": "Empty filename"}), 400

    if not file.filename.lower().endswith(".csv"):
        return jsonify({"error": "Only CSV files are allowed"}), 400

    csv_path = os.path.join(TEMP_DIR, "upload.csv")
    tmp_path = os.path.join(TEMP_DIR, "upload.tmp.csv")

    try:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)

        file.save(tmp_path)
        os.replace(tmp_path, csv_path)

        return jsonify({"status": "uploaded", "path": csv_path})
    finally:
        try:
            file.close()
        except Exception:
            pass

        try:
            if hasattr(file, "stream") and file.stream:
                file.stream.close()
        except Exception:
            pass

        force_memory_cleanup()

@app.route("/save-config", methods=["POST"])
def save_config():
    refresh_worker_state()

    with state_lock:
        if is_running:
            return jsonify({
                "error": "Cannot save config while an algorithm is running"
            }), 409

    data = request.get_json()

    if data is None:
        return jsonify({"error": "No JSON provided"}), 400

    config_path = os.path.join(TEMP_DIR, "config.json")

    with open(config_path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=4)

    force_memory_cleanup()
    return jsonify({"status": "config_saved"})

if __name__ == "__main__":
    mp.freeze_support()
    app.run(host="0.0.0.0", port=5000, debug=False, use_reloader=False)