import os
import base64
import hashlib
import json
import time
import uuid
from flask import Flask, request, render_template, send_file, jsonify, redirect, url_for
import qrcode
from cryptography.fernet import Fernet

app = Flask(__name__)

# --- 配置 ---
UPLOAD_FOLDER = 'uploads'
QR_CODE_FOLDER = 'qr_codes'
RECEIVED_FOLDER = 'received_files'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['QR_CODE_FOLDER'] = QR_CODE_FOLDER
app.config['RECEIVED_FOLDER'] = RECEIVED_FOLDER

# 确保目录存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(QR_CODE_FOLDER, exist_ok=True)
os.makedirs(RECEIVED_FOLDER, exist_ok=True)

# 存储传输会话信息 {session_id: {'key': key, 'chunks': [...], 'metadata': {...}}}
transfer_sessions = {}

# --- 辅助函数 ---

def calculate_sha256(file_path):
    """计算文件的 SHA-256 哈希值"""
    sha256_hash = hashlib.sha256()
    with open(file_path, "rb") as f:
        # Read and update hash string value in blocks of 4K
        for byte_block in iter(lambda: f.read(4096), b""):
            sha256_hash.update(byte_block)
    return sha256_hash.hexdigest()

def split_data(data, chunk_size):
    """将字节数据分割成指定大小的块"""
    chunks = []
    for i in range(0, len(data), chunk_size):
        chunks.append(data[i:i + chunk_size])
    return chunks

def generate_qr_code(data_str, filename):
    """生成二维码并保存为文件"""
    qr = qrcode.QRCode(
        version=1,
        error_correction=qrcode.constants.ERROR_CORRECT_L,
        box_size=10,
        border=4,
    )
    qr.add_data(data_str)
    qr.make(fit=True)

    img = qr.make_image(fill_color="black", back_color="white")
    img.save(os.path.join(app.config['QR_CODE_FOLDER'], filename))

# --- 路由 ---

@app.route('/')
def index():
    """主页 - 文件上传"""
    return render_template('index.html')

@app.route('/upload', methods=['POST'])
def upload_file():
    """处理文件上传和二维码生成"""
    if 'files' not in request.files:
        return "No file part", 400
    files = request.files.getlist('files')
    
    if not files or all(f.filename == '' for f in files):
        return "No selected file", 400

    # 获取用户设置
    try:
        chunk_size_kb = int(request.form.get('chunk_size', 100)) # 默认100KB
        display_speed = int(request.form.get('display_speed', 2)) # 默认2张/秒
    except ValueError:
        return "Invalid input for chunk size or display speed", 400

    chunk_size_bytes = chunk_size_kb * 1024
    delay_ms = int(1000 / display_speed) if display_speed > 0 else 500

    # 为本次传输创建唯一会话ID和加密密钥
    session_id = str(uuid.uuid4())
    key = Fernet.generate_key()
    cipher_suite = Fernet(key)
    
    uploaded_files_info = []
    all_chunks_data = [] # 存储所有文件的加密块数据

    # 1. 处理每个上传的文件
    for file in files:
        if file:
            filename = file.filename
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(filepath)
            
            file_size = os.path.getsize(filepath)
            file_hash = calculate_sha256(filepath)
            file_type = file.content_type or 'unknown'

            # 读取文件内容并Base64编码
            with open(filepath, 'rb') as f:
                file_content = f.read()
            encoded_content = base64.b64encode(file_content).decode('utf-8')

            # 准备元数据
            metadata = {
                'filename': filename,
                'size': file_size,
                'type': file_type,
                'sha256': file_hash
            }

            # 将元数据和内容组合，然后加密
            data_to_encrypt = json.dumps({
                'metadata': metadata,
                'content': encoded_content
            }).encode('utf-8')

            encrypted_data = cipher_suite.encrypt(data_to_encrypt)
            
            # 分割加密后的数据
            chunks = split_data(encrypted_data, chunk_size_bytes)
            all_chunks_data.extend([(chunk, filename) for chunk in chunks]) # 保存块和文件名关联

            uploaded_files_info.append({
                'name': filename,
                'size': file_size,
                'hash': file_hash
            })

    # 2. 为所有块创建二维码
    qr_filenames = []
    total_chunks = len(all_chunks_data)
    for i, (chunk_data, filename) in enumerate(all_chunks_data):
        # 构造二维码数据：会话ID, 块索引, 总块数, 文件名标识, 加密块数据
        # 为了简化，我们将密钥也编码进去（第一个二维码）
        qr_data_dict = {
            'session_id': session_id,
            'index': i,
            'total': total_chunks,
            'filename': filename,
            'data': base64.b64encode(chunk_data).decode('utf-8')
        }
        
        # 第一个二维码包含解密密钥
        if i == 0:
            qr_data_dict['key'] = key.decode('utf-8')

        qr_data_str = json.dumps(qr_data_dict)
        qr_filename = f"qr_{session_id}_{i}.png"
        generate_qr_code(qr_data_str, qr_filename)
        qr_filenames.append(qr_filename)


    # 3. 存储会话信息 (虽然本次简化了密钥分发，但保留结构)
    transfer_sessions[session_id] = {
        'key': key.decode('utf-8'), # 存储为字符串
        'total_chunks': total_chunks,
        'chunks_received': {}, # {index: data}
        'files_info': uploaded_files_info
    }

    # 4. 渲染展示页面
    return render_template('index.html', 
                           qr_files=qr_filenames, 
                           delay_ms=delay_ms,
                           session_id=session_id,
                           files_info=uploaded_files_info)


# --- 接收端模拟 (简化版，实际需要前端JS扫描) ---

@app.route('/scanner')
def scanner():
    """扫描页面"""
    return render_template('scanner.html')

# 模拟接收端扫描并重组文件的API端点
@app.route('/receive_chunk', methods=['POST'])
def receive_chunk():
    """接收扫描到的二维码数据块 (模拟)"""
    data = request.json
    qr_data_str = data.get('qr_data')
    
    if not qr_data_str:
        return jsonify({'status': 'error', 'message': 'No QR data received'}), 400

    try:
        qr_data = json.loads(qr_data_str)
    except json.JSONDecodeError:
        return jsonify({'status': 'error', 'message': 'Invalid QR data format'}), 400

    session_id = qr_data.get('session_id')
    index = qr_data.get('index')
    total = qr_data.get('total')
    filename = qr_data.get('filename')
    encrypted_chunk_b64 = qr_data.get('data')
    key_str = qr_data.get('key') # 可能来自第一个二维码

    if not session_id or index is None or total is None or not encrypted_chunk_b64:
        return jsonify({'status': 'error', 'message': 'Missing data in QR code'}), 400

    # 检查或创建会话
    if session_id not in transfer_sessions:
        transfer_sessions[session_id] = {
            'key': key_str, # 从第一个二维码获取
            'total_chunks': total,
            'chunks_received': {},
            'files_info': {} # 简化处理
        }
    elif key_str:
         transfer_sessions[session_id]['key'] = key_str # 更新密钥（如果第一个二维码后扫到）

    session = transfer_sessions[session_id]
    
    # 解密块
    try:
        cipher_suite = Fernet(session['key'].encode('utf-8'))
        encrypted_chunk = base64.b64decode(encrypted_chunk_b64)
        decrypted_chunk = cipher_suite.decrypt(encrypted_chunk)
        
        # 解析解密后的内容
        chunk_data = json.loads(decrypted_chunk.decode('utf-8'))
        # 这里可以进一步处理 chunk_data['metadata'] 和 chunk_data['content']
        
        # 存储解密后的内容（或原始加密块，取决于重组策略）
        session['chunks_received'][index] = {
            'data': chunk_data, # 存储解密后的完整数据
            'filename': filename
        }
        
    except Exception as e:
        return jsonify({'status': 'error', 'message': f'Decryption failed: {str(e)}'}), 500

    # 检查是否所有块都已接收
    if len(session['chunks_received']) == session['total_chunks']:
        # 重组文件
        try:
            # 按索引排序
            sorted_chunks = sorted(session['chunks_received'].items())
            
            # 假设只有一个文件，或按filename分组处理
            # 这里简化为处理第一个文件
            if sorted_chunks:
                # 重新组装内容 (需要按文件名分组)
                files_content = {}
                for idx, chunk_info in sorted_chunks:
                    fname = chunk_info['filename']
                    content = chunk_info['data']['content'] # Base64内容
                    if fname not in files_content:
                        files_content[fname] = []
                    files_content[fname].append(content)
                
                for fname, parts in files_content.items():
                    # Base64解码所有部分并合并
                    full_base64_content = "".join(parts)
                    file_content = base64.b64decode(full_base64_content)
                    
                    # 保存文件
                    output_path = os.path.join(app.config['RECEIVED_FOLDER'], f"received_{fname}")
                    with open(output_path, 'wb') as f:
                        f.write(file_content)
                    
                    # 可以在这里验证SHA256 (chunk_info['data']['metadata']['sha256'])
                
                # 清理会话
                del transfer_sessions[session_id]
                return jsonify({'status': 'success', 'message': 'File(s) received and saved.', 'download_url': url_for('download_file', filename=f"received_{list(files_content.keys())[0]}")})
            else:
                return jsonify({'status': 'error', 'message': 'No chunks to assemble'}), 500
        except Exception as e:
            return jsonify({'status': 'error', 'message': f'File assembly failed: {str(e)}'}), 500

    return jsonify({'status': 'success', 'message': f'Chunk {index+1}/{total} received'})


@app.route('/download/<filename>')
def download_file(filename):
    """提供下载重组后的文件"""
    try:
        return send_file(os.path.join(app.config['RECEIVED_FOLDER'], filename), as_attachment=True)
    except FileNotFoundError:
        return "File not found", 404


if __name__ == '__main__':
    app.run(debug=True)