from flask import Flask, render_template, request, redirect, url_for, jsonify, send_file, make_response
import os
import json
import time
import hmac
import hashlib
import requests
from datetime import datetime

# ========== 配置项（替换成你的 HF 信息） ==========
HF_TOKEN = "你的HF Access Token（write权限）"
HF_USERNAME = "你的HF用户名（如WACATW）"
HF_REPO = "你的HF私有仓库名（如cat-website-pan）"
HF_API_URL = f"https://huggingface.co/api/repos/{HF_USERNAME}/{HF_REPO}"

# ========== Flask 初始化 ==========
app = Flask(__name__, static_folder='static', static_url_path='')
app.secret_key = "random_secret_key_123"  # 随便填，用于会话加密

# ========== 模拟用户数据（简单版，可替换成 HF 存储） ==========
USERS_FILE = "users.json"
if not os.path.exists(USERS_FILE):
    with open(USERS_FILE, 'w', encoding='utf-8') as f:
        json.dump({}, f)

# ========== HF 文件操作核心函数（替代本地存储） ==========
def hf_upload_file(filename, file_content):
    """上传文件到 HF 仓库"""
    url = f"{HF_API_URL}/files/main/{filename}"
    headers = {
        "Authorization": f"Bearer {HF_TOKEN}",
        "Content-Type": "application/json"
    }
    data = {
        "message": f"Upload file: {filename}",
        "content": file_content,
        "branch": "main"
    }
    response = requests.put(url, headers=headers, json=data)
    return response.ok

def hf_list_files():
    """列出 HF 仓库里的文件"""
    url = f"{HF_API_URL}/tree/main/user-files"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    response = requests.get(url, headers=headers)
    if response.ok:
        return response.json().get("tree", [])
    return []

def hf_download_file(filename):
    """下载 HF 仓库里的文件"""
    url = f"https://huggingface.co/{HF_USERNAME}/{HF_REPO}/raw/main/{filename}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    response = requests.get(url, headers=headers)
    if response.ok:
        return response.content
    return None

def hf_delete_file(filename):
    """删除 HF 仓库里的文件"""
    # 先获取文件的 sha
    url = f"{HF_API_URL}/file/main/{filename}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    response = requests.get(url, headers=headers)
    if not response.ok:
        return False
    sha = response.json().get("sha")
    
    # 删除文件
    delete_url = f"{HF_API_URL}/files/{filename}"
    data = {
        "message": f"Delete file: {filename}",
        "sha": sha,
        "branch": "main"
    }
    delete_resp = requests.delete(delete_url, headers=headers, json=data)
    return delete_resp.ok

# ========== 路由：登录/注册/文件管理 ==========
@app.route('/')
def index():
    """主页，未登录跳登录页"""
    username = request.cookies.get('username')
    if not username:
        return redirect(url_for('login'))
    return redirect(url_for('dashboard'))

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        
        # 验证用户
        with open(USERS_FILE, 'r', encoding='utf-8') as f:
            users = json.load(f)
        if username not in users:
            return render_template('login.html', error="用户名不存在")
        if not hmac.compare_digest(users[username], hashlib.sha256(password.encode()).hexdigest()):
            return render_template('login.html', error="密码错误")
        
        # 登录成功，设置Cookie
        resp = make_response(redirect(url_for('dashboard')))
        resp.set_cookie('username', username, max_age=3600*24)
        return resp
    return render_template('login.html')

@app.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        
        # 验证格式
        if len(username) < 3 or len(username) > 20:
            return render_template('register.html', error="用户名长度3-20字符")
        if len(password) < 6:
            return render_template('register.html', error="密码至少6位")
        
        # 保存用户
        with open(USERS_FILE, 'r+', encoding='utf-8') as f:
            users = json.load(f)
            if username in users:
                return render_template('register.html', error="用户名已存在")
            users[username] = hashlib.sha256(password.encode()).hexdigest()
            f.seek(0)
            json.dump(users, f, ensure_ascii=False)
        
        return redirect(url_for('login'))
    return render_template('register.html')

@app.route('/dashboard')
def dashboard():
    """文件管理页"""
    username = request.cookies.get('username')
    if not username:
        return redirect(url_for('login'))
    # 获取用户的文件列表
    files = hf_list_files()
    user_files = [f for f in files if f['path'].startswith(f'user-files/{username}/')]
    return render_template('dashboard.html', username=username, files=user_files)

@app.route('/upload', methods=['POST'])
def upload():
    """上传文件"""
    username = request.cookies.get('username')
    if not username:
        return "未登录", 401
    
    file = request.files.get('file')
    if not file:
        return "请选择文件", 400
    
    # 生成唯一文件名
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    filename = f"user-files/{username}/{timestamp}_{file.filename}"
    
    # 读取文件内容并上传到 HF（兼容二进制文件）
    file_content = file.read().decode('latin-1')  # 用latin-1兼容所有文件类型
    if hf_upload_file(filename, file_content):
        return redirect(url_for('dashboard'))
    return "上传失败", 500

@app.route('/download/<path:filename>')
def download(filename):
    """下载文件"""
    username = request.cookies.get('username')
    if not username or not filename.startswith(f'user-files/{username}/'):
        return "无权限", 403
    
    # 从 HF 下载文件
    content = hf_download_file(filename)
    if not content:
        return "文件不存在", 404
    
    # 发送文件给用户
    resp = make_response(content)
    resp.headers['Content-Disposition'] = f'attachment; filename="{os.path.basename(filename)}"'
    return resp

@app.route('/delete/<path:filename>')
def delete(filename):
    """删除文件"""
    username = request.cookies.get('username')
    if not username or not filename.startswith(f'user-files/{username}/'):
        return "无权限", 403
    
    if hf_delete_file(filename):
        return redirect(url_for('dashboard'))
    return "删除失败", 500

@app.route('/logout')
def logout():
    """退出登录"""
    resp = make_response(redirect(url_for('login')))
    resp.delete_cookie('username')
    return resp

# ========== 启动应用 ==========
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)  # HF Spaces 默认端口 7860