# backup_utils.py - 数据备份工具
"""管理员数据备份：导出所有表数据为 JSON，支持打包上传文件"""
import json
import os
import zipfile
from datetime import datetime
from models import beijing_now
from io import BytesIO

from models import (
    db, User, Exam, Submission, Draft, Contest, ContestMembership,
    ContestApplication, Post, Reply, Poll, Report, Bookmark, Reaction,
    Notification, EditHistory, ContestRegistration, TeacherApplication,
    Friend, ExamBookmark, ChatRoom, ChatRoomMember, Message, MessageReaction,
    QuestionBankItem, InviteCode, SystemNotification
)


def _serialize_value(val):
    """序列化单个值为 JSON 可存储格式"""
    if val is None:
        return None
    if isinstance(val, datetime):
        return val.strftime('%Y-%m-%d %H:%M:%S')
    if isinstance(val, bool):
        return val
    if isinstance(val, (int, float, str)):
        return val
    if isinstance(val, (list, dict)):
        return val
    return str(val)


def _model_to_dict(obj):
    """将 SQLAlchemy 模型实例转为可序列化字典"""
    if obj is None:
        return None
    d = {}
    for col in obj.__table__.columns:
        val = getattr(obj, col.name)
        d[col.name] = _serialize_value(val)
    return d


# 按依赖顺序排列的表（父表在前）
BACKUP_MODELS = [
    ('user', User),
    ('contest', Contest),
    ('contest_membership', ContestMembership),
    ('contest_application', ContestApplication),
    ('contest_registration', ContestRegistration),
    ('question_bank_item', QuestionBankItem),
    ('exam', Exam),
    ('submission', Submission),
    ('draft', Draft),
    ('post', Post),
    ('reply', Reply),
    ('poll', Poll),
    ('reaction', Reaction),
    ('bookmark', Bookmark),
    ('report', Report),
    ('notification', Notification),
    ('edit_history', EditHistory),
    ('teacher_application', TeacherApplication),
    ('friend', Friend),
    ('exam_bookmark', ExamBookmark),
    ('chat_room', ChatRoom),
    ('chat_room_member', ChatRoomMember),
    ('message', Message),
    ('message_reaction', MessageReaction),
    ('invite_code', InviteCode),
    ('system_notification', SystemNotification),
]


def export_all_data():
    """导出所有表数据为 JSON 字典"""
    data = {
        'meta': {
            'exported_at': beijing_now().strftime('%Y-%m-%d %H:%M:%S'),
            'version': '1.0',
        },
        'tables': {}
    }
    for table_name, model in BACKUP_MODELS:
        try:
            rows = model.query.order_by(model.id).all()
            data['tables'][table_name] = [_model_to_dict(r) for r in rows]
        except Exception as e:
            data['tables'][table_name] = {'_error': str(e)}
    return data


def create_backup_zip(include_uploads=True, upload_folder=None):
    """
    创建完整备份 ZIP 包
    :param include_uploads: 是否包含上传文件
    :param upload_folder: 上传目录路径，默认 static/uploads
    :return: BytesIO 中的 ZIP 文件
    """
    buf = BytesIO()
    with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
        # 1. 导出数据 JSON
        data = export_all_data()
        zf.writestr('data.json', json.dumps(data, ensure_ascii=False, indent=2))

        # 2. 可选：上传文件
        if include_uploads and upload_folder and os.path.isdir(upload_folder):
            for root, dirs, files in os.walk(upload_folder):
                for f in files:
                    fp = os.path.join(root, f)
                    arcname = 'uploads/' + os.path.relpath(fp, upload_folder).replace('\\', '/')
                    zf.write(fp, arcname, compress_type=zipfile.ZIP_DEFLATED)

    buf.seek(0)
    return buf