from __future__ import annotations import os from datetime import datetime from flask import ( Blueprint, abort, current_app, flash, jsonify, redirect, render_template, request, send_file, session, url_for, ) from flask_login import current_user, login_required, login_user, logout_user from .config import BASE_DIR from .extensions import db from .models import UploadRecord, User from .prediction import PredictionError, run_prediction from .security import new_captcha bp = Blueprint("main", __name__) @bp.route("/") def index(): if current_user.is_authenticated: return redirect(url_for("main.home")) return redirect(url_for("main.login")) @bp.route("/login", methods=["GET", "POST"]) def login(): if current_user.is_authenticated: return redirect(url_for("main.home")) if request.method == "GET": session["captcha"] = new_captcha() return render_template("login.html", mode="login", captcha=session["captcha"]) username = request.form.get("username", "").strip() password = request.form.get("password", "") captcha_input = request.form.get("captcha", "").strip().upper() if captcha_input != session.get("captcha", ""): flash("验证码错误", "error") session["captcha"] = new_captcha() return render_template("login.html", mode="login", captcha=session["captcha"]), 400 user = User.query.filter_by(username=username).first() if not user or not user.check_password(password): flash("用户名或密码错误", "error") session["captcha"] = new_captcha() return render_template("login.html", mode="login", captcha=session["captcha"]), 400 login_user(user, remember=bool(request.form.get("remember"))) return redirect(url_for("main.home")) @bp.route("/register", methods=["GET", "POST"]) def register(): if not current_app.config["ALLOW_REGISTRATION"]: abort(404) if request.method == "GET": return render_template("login.html", mode="register", captcha="") username = request.form.get("username", "").strip() password = request.form.get("password", "") if not username: flash("用户名不能为空", "error") return render_template("login.html", mode="register", captcha=""), 400 if len(password) < 6: flash("密码至少需要 6 位", "error") return render_template("login.html", mode="register", captcha=""), 400 if User.query.filter_by(username=username).first(): flash("用户名已存在", "error") return render_template("login.html", mode="register", captcha=""), 400 user = User(username=username, is_admin=False) user.set_password(password) db.session.add(user) db.session.commit() flash("注册成功，请登录", "info") session["captcha"] = new_captcha() return render_template("login.html", mode="login", captcha=session["captcha"]) @bp.route("/logout", methods=["POST"]) @login_required def logout(): logout_user() return redirect(url_for("main.login")) @bp.route("/home") @login_required def home(): return render_template("home.html") @bp.route("/history") @login_required def history_page(): records = ( UploadRecord.query.filter_by(user_id=current_user.id) .order_by(UploadRecord.upload_time.desc()) .all() ) return render_template("history.html", records=records) @bp.route("/admin") @login_required def admin_dashboard(): if not current_user.is_admin: abort(403) records = UploadRecord.query.order_by(UploadRecord.upload_time.desc()).all() return render_template("admin.html", records=records) @bp.route("/download//") @login_required def download_file(record_id: int, file_type: str): record = db.session.get(UploadRecord, record_id) if record is None: abort(404) if not (current_user.is_admin or current_user.id == record.user_id): abort(403) if file_type == "original": file_path = record.saved_path elif file_type == "prediction": file_path = record.prediction_path else: abort(404) if not os.path.exists(file_path): abort(404) return send_file(file_path, as_attachment=True) @bp.route("/download_template") def download_template(): template_path = BASE_DIR / "example.xlsx" if not template_path.exists(): abort(404) return send_file(template_path, as_attachment=True, download_name="example.xlsx") @bp.route("/result") @login_required def result_page(): result = session.get("last_result") return render_template("result.html", result=result) @bp.route("/predict", methods=["POST"]) @login_required def predict(): model = current_app.config.get("RSF_MODEL") if model is None: return jsonify({"error": "模型未成功加载，请检查模型文件。"}), 500 uploaded = request.files.get("file") if uploaded is None or uploaded.filename == "": return jsonify({"error": "未选择文件"}), 400 try: artifacts = run_prediction(uploaded, int(current_user.id), model) except PredictionError as exc: return jsonify({"error": exc.message}), exc.status_code record = UploadRecord( user_id=current_user.id, original_filename=artifacts.original_filename, saved_path=str(artifacts.saved_path), prediction_path=str(artifacts.excel_path), image_path=str(artifacts.image_path), ) db.session.add(record) db.session.commit() last_result = { "original_filename": artifacts.original_filename, "generated_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "image_url": url_for("static", filename=f"images/{artifacts.image_filename}"), "importance_url": ( url_for("static", filename=f"images/{artifacts.importance_filename}") if artifacts.importance_filename else None ), "excel_url": url_for("main.download_file", record_id=record.id, file_type="prediction"), "result_url": url_for("main.result_page"), "sample_count": int(artifacts.sample_count), "summary_rows": artifacts.summary_rows[:3], "analysis_text": artifacts.analysis_text, } session["last_result"] = last_result return jsonify( { "message": "预测成功", "image_url": last_result["image_url"], "importance_url": last_result["importance_url"], "excel_url": last_result["excel_url"], "result_url": last_result["result_url"], "sample_count": last_result["sample_count"], "original_filename": artifacts.original_filename, } )