from flask import Flask, send_file, jsonify, request
import zipfile
import os
import tempfile
import threading
from datetime import datetime

app = Flask(__name__)

class ZipCreator:
    def __init__(self, source_path, temp_dir):
        self.source_path = f"/var/www/{source_path}"
        self.temp_dir = temp_dir
        #self.zip_file_path = os.path.join(temp_dir, 'files.zip')
        self.zip_file_name = os.path.basename(source_path) + '.zip'
        self.zip_file_path = os.path.join(temp_dir, self.zip_file_name)
        self.status = "pending"
        self.error = None

    def create_zip(self):
        try:
            with zipfile.ZipFile(self.zip_file_path, 'w') as zip_file:
                for root, dirs, files in os.walk(self.source_path):
                    # Clear the dirs list to prevent descending into subdirectories
                    dirs.clear()
                    for file in files:
                        if file.lower().endswith(('.jpg', '.jpeg')): # Filter only JPEG files
                            file_path = os.path.join(root, file)
                            zip_file.write(file_path, os.path.relpath(file_path, self.source_path))
            self.status = "completed"
        except Exception as e:
            self.status = "failed"
            self.error = str(e)
        finally:
            # Ensure the temp directory is cleaned up after download
            def cleanup_temp_dir():
                import time
                time.sleep(30)  # Wait 30 seconds to ensure download is complete
                try:
                    os.remove(self.zip_file_path)
                    os.rmdir(self.temp_dir)
                except Exception as e:
                    print(f"Error cleaning up temp directory: {e}")

            threading.Thread(target=cleanup_temp_dir).start()

# Global variable to hold the current zip creation task
current_task = None

@app.route("/start_zip", methods=["POST"])
def start_zip():
    global current_task
    source_path = request.json['source_path']
    temp_dir = tempfile.mkdtemp()
    current_task = ZipCreator(source_path, temp_dir)
    threading.Thread(target=current_task.create_zip).start()
    return jsonify({"message": "Zip creation started"})

@app.route("/status", methods=["GET"])
def status():
    global current_task
    if current_task:
        return jsonify({"status": current_task.status, "error": current_task.error})
    else:
        return jsonify({"status": "no_task"})

@app.route("/download", methods=["GET"])
def download():
    global current_task
    if current_task and current_task.status == "completed":
        return send_file(current_task.zip_file_path, as_attachment=True)
    elif current_task and current_task.status == "failed":
        return jsonify({"error": current_task.error}), 500
    else:
        return jsonify({"error": "Zip not ready"}), 400

if __name__ == "__main__":
    app.run(host='0.0.0.0', port=5000)