from flask import Flask, render_template, request, flash, send_from_directory, redirect, url_for, jsonify
from flask_wtf import FlaskForm
from wtforms import StringField, SelectField, SubmitField
from wtforms.validators import DataRequired
import mysql.connector
from PIL import Image
import os
import shutil
import hashlib
from math import ceil

app = Flask(__name__)
app.config['SECRET_KEY'] = 'linux'

# 数据库
db_config = {
    'host': '192.168.1.245',
    'user': 'api',
    'password': '35NR6idXCHkQdx4M',
    'database': 'api'
}

def get_db_connection():
    return mysql.connector.connect(**db_config)

class TagImageForm(FlaskForm):
    filename = StringField('New Filename', validators=[DataRequired()])
    theme = SelectField('Theme', choices=[
        ('light', 'Light'),
        ('dark', 'Dark'),
        ('favicon', 'Favicon'),
        ('rainyun', 'Rainyun'),
        ('fox', 'Fox'),
        ('bj', 'BJ')
    ], validators=[DataRequired()])
    et = SelectField('ET', choices=[
        ('pc', 'PC'),
        ('phone', 'Phone'),
        ('all', 'All')
    ], validators=[DataRequired()], default='pc')
    submit = SubmitField('Tag and Move Image')

@app.route('/', methods=['GET', 'POST'])
def index():
    form = TagImageForm()
    page = request.args.get('page', 1, type=int)
    per_page = 6

    if form.validate_on_submit():
        new_filename = form.filename.data
        theme = form.theme.data
        et = form.et.data
        original_filename = request.form.get('original_filename')

        try:
            conn = get_db_connection()
            cursor = conn.cursor()
            query = "INSERT INTO images (filename, theme, et) VALUES (%s, %s, %s)"
            cursor.execute(query, (new_filename, theme, et))
            conn.commit()

            catch_dir = os.path.join(app.root_path, 'catch')
            img_dir = os.path.join(app.root_path, 'img')
            current_path = os.path.join(catch_dir, original_filename)
            new_path = os.path.join(img_dir, new_filename)
            shutil.move(current_path, new_path)

            
            conn.commit()

            flash('图片已标记成功', 'success')
            return redirect(url_for('index'))

        except Exception as e:
            flash(f'Error: {str(e)}', 'danger')
            return redirect(url_for('index'))

    catch_dir = os.path.join(app.root_path, 'catch')
    all_images = [f for f in os.listdir(catch_dir) if os.path.isfile(os.path.join(catch_dir, f)) and f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.ico', '.webp'))]

    total_pages = ceil(len(all_images) / per_page)
    start = (page - 1) * per_page
    end = start + per_page
    images = all_images[start:end]

    return render_template('index.html', form=form, images=images, page=page, total_pages=total_pages)

@app.route('/delete/<filename>', methods=['GET'])
def delete_image(filename):
    try:
        # 删除图片
        catch_dir = os.path.join(app.root_path, 'catch')
        file_path = os.path.join(catch_dir, filename)
        if os.path.exists(file_path):
            os.remove(file_path)
            flash('图片删除成功', 'success')
        else:
            flash('未找到图片', 'warning')
    except Exception as e:
        flash(f'Error: {str(e)}', 'danger')
    
    return redirect(url_for('index'))

@app.route('/catch/<path:filename>')
def serve_catch_file(filename):
    catch_dir = os.path.join(app.root_path, 'catch')
    return send_from_directory(catch_dir, filename)

@app.route('/img/<path:filename>')
def serve_img_file(filename):
    img_dir = os.path.join(app.root_path, 'img')
    return send_from_directory(img_dir, filename)

# 短哈希
@app.route('/generate-short-hash', methods=['POST'])
def generate_short_hash():
    filename = request.form.get('filename')
    short_hash = generate_short_hash(filename)
    return jsonify({'short_hash': short_hash})


Image.MAX_IMAGE_PIXELS = 1000000000

def generate_short_hash(filename, length=8):
    hash_object = hashlib.sha256(filename.encode())
    hex_dig = hash_object.hexdigest()
    return hex_dig[:length]

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