slideshow/Slideshow.py

import json
import os
import subprocess
import sys
import threading
import time
import uuid
from flask import Flask, render_template, request, redirect, url_for, Response, send_from_directory
from werkzeug.middleware.proxy_fix import ProxyFix

from db import init_db, search_images, get_image_count, reset_db

app = Flask(__name__, static_folder='Pictures', static_url_path='/pictures')
app.wsgi_app = ProxyFix(app.wsgi_app, x_prefix=1)
init_db()

# job_id -> {'done': int, 'total': int, 'finished': bool, 'site': str, 'tags': str}
downloads = {}


@app.route('/previews/<path:filename>')
def serve_preview(filename):
    return send_from_directory('Previews', filename)


@app.route('/')
def slideshow():
    raw_query = request.args.get('tags', '').strip()
    selected_site = request.args.get('_site', 'e621')
    sort = request.args.get('_sort', '')
    results = search_images(raw_query, sort=sort or None)
    base_dir = os.path.dirname(os.path.abspath(__file__))
    pictures_dir = os.path.join(base_dir, 'Pictures')
    previews_dir = os.path.join(base_dir, 'Previews')

    image_urls    = []
    preview_urls  = []
    post_urls     = []
    tags_list     = []
    file_sizes    = []
    preview_sizes = []
    created_ats   = []

    for r in results:
        fname = r['filename']
        pname = r['preview_filename']
        image_urls.append(f'pictures/{fname}')
        post_urls.append(r['post_url'])
        tags_list.append(r['tags'].split())
        file_sizes.append(os.path.getsize(os.path.join(pictures_dir, fname)))
        created_ats.append(r['created_at'])

        ppath = os.path.join(previews_dir, pname) if pname else ''
        if pname and os.path.exists(ppath):
            preview_urls.append(f'previews/{pname}')
            preview_sizes.append(os.path.getsize(ppath))
        else:
            preview_urls.append(f'pictures/{fname}')
            preview_sizes.append(file_sizes[-1])

    active_tags = raw_query.split() if raw_query else []
    job_id = request.args.get('job_id')
    return render_template(
        'slideshow.html',
        images=image_urls,
        preview_images=preview_urls,
        post_urls=post_urls,
        tags_list=tags_list,
        file_sizes=file_sizes,
        preview_sizes=preview_sizes,
        created_ats=created_ats,
        active_tags=active_tags,
        tag_query=raw_query,
        selected_site=selected_site,
        sort=sort,
        job_id=job_id,
    )


@app.route('/admin')
def admin():
    base_dir = os.path.dirname(os.path.abspath(__file__))
    pictures_dir = os.path.join(base_dir, 'Pictures')
    previews_dir = os.path.join(base_dir, 'Previews')

    def folder_size(path):
        total = 0
        for dirpath, _, filenames in os.walk(path):
            for f in filenames:
                fp = os.path.join(dirpath, f)
                if not os.path.islink(fp):
                    total += os.path.getsize(fp)
        return total

    def fmt_size(n):
        if n >= 1_073_741_824:
            return f'{n / 1_073_741_824:.2f} GB'
        return f'{n / 1_048_576:.1f} MB'

    return render_template('admin.html',
        pic_count=get_image_count(),
        pics_size=fmt_size(folder_size(pictures_dir)),
        previews_size=fmt_size(folder_size(previews_dir)),
    )


@app.route('/admin/reset', methods=['POST'])
def admin_reset():
    base_dir = os.path.dirname(os.path.abspath(__file__))
    for folder in ('Pictures', 'Previews'):
        folder_path = os.path.join(base_dir, folder)
        for dirpath, _, filenames in os.walk(folder_path):
            for f in filenames:
                if f != '.gitkeep':
                    os.remove(os.path.join(dirpath, f))
    reset_db()
    return redirect(url_for('admin'))


@app.route('/download', methods=['POST'])
def download():
    tags = request.form.get('tags', '').strip()
    site = request.form.get('site', 'e621')
    if not tags:
        return redirect(url_for('slideshow', tags=tags))

    job_id = uuid.uuid4().hex[:8]
    downloads[job_id] = {'done': 0, 'total': 0, 'finished': False, 'site': site, 'tags': tags, 'status': ''}

    proc = subprocess.Popen(
        [sys.executable, 'downloader.py', '--site', site, '--query', tags],
        cwd=app.root_path,
        stdout=subprocess.PIPE,
        text=True,
    )

    def read_stdout():
        for line in proc.stdout:
            line = line.strip()
            if line.startswith('status:'):
                downloads[job_id]['status'] = line[7:]
            elif line.startswith('total:'):
                downloads[job_id]['total'] = int(line.split(':')[1])
            elif line.startswith('progress:'):
                done, total = line.split(':')[1].split('/')
                downloads[job_id]['done'] = int(done)
                downloads[job_id]['total'] = int(total)
            elif line == 'done':
                downloads[job_id]['finished'] = True
        downloads[job_id]['finished'] = True

    threading.Thread(target=read_stdout, daemon=True).start()
    return redirect(url_for('slideshow', tags=tags, job_id=job_id))


@app.route('/download/progress/<job_id>')
def download_progress(job_id):
    def generate():
        while True:
            info = downloads.get(job_id)
            if not info:
                yield f'data: {json.dumps({"error": "not found"})}\n\n'
                break
            yield f'data: {json.dumps({"done": info["done"], "total": info["total"], "finished": info["finished"], "status": info["status"]})}\n\n'
            if info['finished']:
                break
            time.sleep(0.3)

    return Response(
        generate(),
        mimetype='text/event-stream',
        headers={'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'},
    )


if __name__ == '__main__':
    app.run(host='0.0.0.0', debug=False)