import os
import smtplib
import yaml
import psutil
import threading
import time
from flask import Flask, render_template
from email.mime.text import MIMEText
import subprocess

app = Flask(__name__)

# Load configuration
try:
    with open("config.yaml") as f:
        config = yaml.safe_load(f)
        print("[DEBUG] Loaded config.yaml successfully.")
except Exception as e:
    print(f"[ERROR] Failed to load config.yaml: {e}")
    exit(1)

results = {}

def get_folder_size(path):
    print(f"[DEBUG] Calculating size for: {path}")
    try:
        out = subprocess.check_output(["du", "-sb", path], timeout=60)
        size_bytes = int(out.decode().split()[0])
        print(f"[DEBUG] Total size for {path}: {size_bytes / 1e9:.2f} GB")
        return size_bytes
    except Exception as e:
        print(f"[ERROR] Failed to get size for {path}: {e}")
        return 0

def send_alert(folder, used, limit):
    print(f"[ALERT] Sending email for {folder}: used={used}, limit={limit}")
    try:
        msg = MIMEText(f"Folder {folder} is over limit: {used / 1e9:.2f} GB used of {limit / 1e9:.2f} GB")
        msg['Subject'] = f"[ALERT] Disk usage for {folder}"
        msg['From'] = config['email']['from']
        msg['To'] = config['email']['to']

        with smtplib.SMTP(config['email']['smtp_server'], config['email']['smtp_port']) as s:
            s.starttls()
            s.login(config['email']['username'], config['email']['password'])
            s.send_message(msg)
        print(f"[INFO] Email alert sent for {folder}")
    except Exception as e:
        print(f"[ERROR] Failed to send email alert: {e}")

def check_folders():
    print("[INFO] Running folder usage check...")
    new_results = {}
    for folder, limit_str in config['folders'].items():
        if not os.path.exists(folder):
            print(f"[WARN] Folder does not exist: {folder}")
            continue
        try:
            limit = int(limit_str.replace("GB", "").strip()) * 1024 ** 3
            used = get_folder_size(folder)
            exceeded = used > limit
            new_results[folder] = {
                "used_bytes": used,
                "limit_bytes": limit,
                "used_gb": used / (1024 ** 3),
                "limit_gb": limit / (1024 ** 3),
                "exceeded": exceeded
            }
            if exceeded:
                send_alert(folder, used, limit)
        except Exception as e:
            print(f"[ERROR] Failed to check folder {folder}: {e}")
    results.clear()
    results.update(new_results)

def schedule_check(interval_sec):
    def loop():
        while True:
            check_folders()
            time.sleep(interval_sec)
    t = threading.Thread(target=loop, daemon=True)
    t.start()

@app.route("/")
def index():
    import datetime
    return render_template("index.html", results=results, now=datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'))


if __name__ == "__main__":
    interval = config.get("schedule", {}).get("seconds", 86400)  # default 5 minutes
    print(f"[INFO] Starting scheduled checks every {interval} seconds.")
    schedule_check(interval)

    if config.get("webui", {}).get("enabled", False):
        port = config["webui"].get("port", 8080)
        print(f"[INFO] Starting Flask Web UI on port {port}")
        app.run(host="0.0.0.0", port=port)
    else:
        print("[INFO] Web UI disabled or missing from config.")