from flask import Flask, render_template, request, redirect, url_for, send_from_directory, session, jsonify
from flask_socketio import SocketIO, emit, disconnect
import os
import json
from werkzeug.security import check_password_hash
from PIL import Image, UnidentifiedImageError
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler

app = Flask(__name__)

# Load configuration from keys.json
with open('db/keys.json', 'r') as f:
    config = json.load(f)

app.config['SECRET_KEY'] = config['key']
socketio = SocketIO(app)

app.config['UPLOAD_FOLDER'] = 'images'  # Folder to store images
app.config['DUMP_FOLDER'] = 'dump_images'
app.config['CACHE_FOLDER'] = 'cache'  # Folder to store cached resized images
app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg', 'gif'}

USERNAME = config['login_info']['username']
PASSWORD_HASH = config['login_info']['password_hash']
SALT = config['login_info']['salt']

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']

def is_image_valid(image_path):
    try:
        with Image.open(image_path) as img:
            img.verify()  # Verify that it is, in fact, an image
        return True
    except Exception:
        return False

def resize_image(image_path, cache_path, size=(800, 450)):
    with Image.open(image_path) as img:
        img.thumbnail(size)
        if img.mode in ("RGBA", "P"):  # Convert to RGB if necessary
            img = img.convert("RGB")
        img.save(cache_path, format="JPEG")

def create_cache():
    images = os.listdir(app.config['UPLOAD_FOLDER'])
    cached_images = []

    if not os.path.exists(app.config['CACHE_FOLDER']):
        os.makedirs(app.config['CACHE_FOLDER'])

    for image in images:
        image_path = os.path.join(app.config['UPLOAD_FOLDER'], image)
        cache_path = os.path.join(app.config['CACHE_FOLDER'], image)

        if not os.path.exists(cache_path) and is_image_valid(image_path):
            resize_image(image_path, cache_path)

        if is_image_valid(image_path):
            cached_images.append(image)
    
    return cached_images

@app.route('/')
def index():
    return redirect(url_for('login'))

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']
        if username == USERNAME and check_password_hash(PASSWORD_HASH, password + SALT):
            session['logged_in'] = True
            return redirect(url_for('gallery'))
        else:
            error = 'Invalid credentials'
            return render_template('login.html', error=error)
    return render_template('login.html')

@app.route('/logout')
def logout():
    session.pop('logged_in', None)
    return redirect(url_for('login'))

@app.route('/gallery')
def gallery():
    if not session.get('logged_in'):
        return redirect(url_for('login'))
    
    cached_images = create_cache()
    batches = []
    current_batch = []
    current_batch_number = None
    unstructured_images = []

    for image in cached_images:
        parts = image.split('_')
        if len(parts) == 2 and parts[1].split('.')[0].isdigit():
            batch_number = int(parts[0])
            if current_batch_number is None:
                current_batch_number = batch_number
            if batch_number != current_batch_number:
                batches.append(current_batch)
                current_batch = []
                current_batch_number = batch_number
            current_batch.append(image)
        else:
            unstructured_images.append(image)

    if current_batch:
        batches.append(current_batch)
    if unstructured_images:
        batches.append(unstructured_images)

    return render_template('gallery.html', batches=batches)

@app.route('/images/<filename>')
def uploaded_file(filename):
    if not session.get('logged_in'):
        return redirect(url_for('login'))
    
    image_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    if is_image_valid(image_path):
        return send_from_directory(app.config['UPLOAD_FOLDER'], filename)
    else:
        return "Invalid image", 404

@app.route('/cached/<filename>')
def cached_file(filename):
    if not session.get('logged_in'):
        return redirect(url_for('login'))
    
    image_path = os.path.join(app.config['CACHE_FOLDER'], filename)
    if is_image_valid(image_path):
        return send_from_directory(app.config['CACHE_FOLDER'], filename)
    else:
        return "Invalid image", 404

@app.route('/api/images')
def api_images():
    if not session.get('logged_in'):
        return jsonify([])

    images = sorted(
        os.listdir(app.config['UPLOAD_FOLDER']),
        key=lambda x: os.path.getctime(os.path.join(app.config['UPLOAD_FOLDER'], x)),
        reverse=True
    )
    valid_images = [img for img in images if is_image_valid(os.path.join(app.config['UPLOAD_FOLDER'], img))]
    return jsonify(valid_images)

@app.route('/dump')
def dump():
    if not session.get('logged_in'):
        return redirect(url_for('login'))
    
    dump_folder = app.config['DUMP_FOLDER'] # Folder to store dump images
    images = sorted(
        os.listdir(dump_folder),
        key=lambda x: os.path.getctime(os.path.join(dump_folder, x)),
        reverse=True
    )
    valid_images = [img for img in images if is_image_valid(os.path.join(dump_folder, img))]
    return render_template('dump.html', images=valid_images)

def emit_gallery_update():
    socketio.emit('update_gallery')

class Watcher(FileSystemEventHandler):
    def on_modified(self, event):
        if not event.is_directory and event.src_path.startswith(app.config['UPLOAD_FOLDER']):
            create_cache()
            emit_gallery_update()

    def on_created(self, event):
        if not event.is_directory and event.src_path.startswith(app.config['UPLOAD_FOLDER']):
            create_cache()
            emit_gallery_update()

    def on_deleted(self, event):
        if not event.is_directory and event.src_path.startswith(app.config['UPLOAD_FOLDER']):
            create_cache()
            emit_gallery_update()

@socketio.on('connect')
def handle_connect():
    if not session.get('logged_in'):
        disconnect()

if __name__ == '__main__':
    observer = Observer()
    observer.schedule(Watcher(), path=app.config['UPLOAD_FOLDER'], recursive=False)
    observer.start()
    try:
        context = ('cert.pem', 'key.pem')  # Replace with your self-signed certificate and key files
        socketio.run(app, debug=True, ssl_context=context)
    except KeyboardInterrupt:
        observer.stop()
    observer.join()