from flask import Flask, request, jsonify, redirect
from flask_cors import CORS
import json
import os
import random
import string
from datetime import datetime, timedelta

app = Flask(__name__)
CORS(app)

DATA_FILE = 'link_data.json'

def load_links():
    """加载链接数据"""
    if not os.path.exists(DATA_FILE):
        return {}
    try:
        with open(DATA_FILE, 'r', encoding='utf-8') as f:
            return json.load(f)
    except:
        return {}

def save_links(links):
    """保存链接数据"""
    with open(DATA_FILE, 'w', encoding='utf-8') as f:
        json.dump(links, f, ensure_ascii=False, indent=2)

def clean_expired_links():
    """清理过期的链接"""
    links = load_links()
    current_time = datetime.now()
    expired_codes = []
    
    for code, data in links.items():
        expires_at = data.get('expires_at')
        if expires_at and expires_at != 'never':
            expire_time = datetime.fromisoformat(expires_at)
            if current_time > expire_time:
                expired_codes.append(code)
    
    for code in expired_codes:
        del links[code]
    
    if expired_codes:
        save_links(links)
    
    return len(expired_codes)

def is_link_expired(link_data):
    """检查链接是否过期"""
    expires_at = link_data.get('expires_at')
    if not expires_at or expires_at == 'never':
        return False
    
    expire_time = datetime.fromisoformat(expires_at)
    return datetime.now() > expire_time

def calculate_expire_time(expire_option):
    """计算过期时间"""
    if expire_option == 'never':
        return 'never'
    
    expire_map = {
        '10min': timedelta(minutes=10),
        '1hour': timedelta(hours=1),
        '1day': timedelta(days=1),
        '3days': timedelta(days=3),
        '1month': timedelta(days=30),
        '1year': timedelta(days=365)
    }
    
    delta = expire_map.get(expire_option)
    if delta:
        return (datetime.now() + delta).isoformat()
    return 'never'

def generate_random_code():
    """生成随机4位字母数字组合（包含大小写）"""
    chars = string.ascii_letters + string.digits  # a-zA-Z0-9
    return ''.join(random.choices(chars, k=4))



#==========================后端API接口==========================#
@app.route('/')
def index():
    """首页"""
    return jsonify({
        'message': '萌芽短链接分发系统',
        "author": "树萌芽",
        "API":
        {
            "/api/check/XXXX":"检查短链接是否已被使用",
            "/api/generate":"生成随机可用的短链接代码",
            "/api/create":"创建短链接",
            "/XXXX":"跳转到目标链接",
            "/api/links":"获取所有链接（可选功能）",
            "/api/redirect/XXXX":"API方式获取跳转链接（供前端路由使用）",
            "":"",

        }
    })

@app.route('/api/check/<code>', methods=['GET'])
def check_code(code):
    """检查短链接是否已被使用"""
    if len(code) != 4:
        return jsonify({'valid': False, 'message': '链接代码必须是4位字符'}), 400
    
    links = load_links()
    exists = code in links
    
    return jsonify({
        'exists': exists,
        'available': not exists
    })

@app.route('/api/generate', methods=['GET'])
def generate_code():
    """生成随机可用的短链接代码"""
    links = load_links()
    max_attempts = 100
    
    for _ in range(max_attempts):
        code = generate_random_code()
        if code not in links:
            return jsonify({
                'code': code,
                'success': True
            })
    
    return jsonify({
        'success': False,
        'message': '生成失败，请重试'
    }), 500

@app.route('/api/create', methods=['POST'])
def create_link():
    """创建短链接"""
    data = request.json
    code = data.get('code', '').strip()
    target_url = data.get('target_url', '').strip()
    expire_option = data.get('expire_option', 'never')
    
    # 验证输入
    if not code or not target_url:
        return jsonify({
            'success': False,
            'message': '链接代码和目标地址不能为空'
        }), 400
    
    if len(code) != 4:
        return jsonify({
            'success': False,
            'message': '链接代码必须是4位字符'
        }), 400
    
    # 清理过期链接
    clean_expired_links()
    
    # 检查是否已存在
    links = load_links()
    if code in links:
        return jsonify({
            'success': False,
            'message': '该链接代码已被使用，请选择其他代码'
        }), 409
    
    # 计算过期时间
    expires_at = calculate_expire_time(expire_option)
    
    # 保存链接
    links[code] = {
        'target_url': target_url,
        'created_at': datetime.now().isoformat(),
        'expires_at': expires_at,
        'expire_option': expire_option,
        'visit_count': 0
    }
    save_links(links)
    
    return jsonify({
        'success': True,
        'code': code,
        'short_url': f'short.shumengya.top/{code}',
        'expires_at': expires_at,
        'message': '短链接创建成功'
    })

@app.route('/<code>', methods=['GET'])
def redirect_link(code):
    """跳转到目标链接"""
    if len(code) != 4:
        return jsonify({
            'error': '无效的短链接',
            'message': '链接代码必须是4位字符'
        }), 404
    
    # 清理过期链接
    clean_expired_links()
    
    links = load_links()
    if code not in links:
        return jsonify({
            'error': '短链接不存在',
            'message': f'未找到代码为 {code} 的短链接'
        }), 404
    
    link_data = links[code]
    
    # 检查是否过期
    if is_link_expired(link_data):
        del links[code]
        save_links(links)
        return jsonify({
            'error': '短链接已过期',
            'message': f'代码为 {code} 的短链接已过期'
        }), 410
    
    # 增加访问次数
    link_data['visit_count'] = link_data.get('visit_count', 0) + 1
    links[code] = link_data
    save_links(links)
    
    target_url = link_data['target_url']
    
    # 确保URL包含协议
    if not target_url.startswith(('http://', 'https://')):
        target_url = 'http://' + target_url
    
    return redirect(target_url, code=302)

@app.route('/api/links', methods=['GET'])
def get_all_links():
    """获取所有链接（可选功能）"""
    # 清理过期链接
    clean_expired_links()
    
    links = load_links()
    return jsonify({
        'success': True,
        'count': len(links),
        'links': links
    })

@app.route('/api/redirect/<code>', methods=['GET'])
def api_redirect(code):
    """API方式获取跳转链接（供前端路由使用）"""
    if len(code) != 4:
        return jsonify({
            'success': False,
            'message': '链接代码必须是4位字符'
        }), 404
    
    # 清理过期链接
    clean_expired_links()
    
    links = load_links()
    if code not in links:
        return jsonify({
            'success': False,
            'message': f'未找到代码为 {code} 的短链接'
        }), 404
    
    link_data = links[code]
    
    # 检查是否过期
    if is_link_expired(link_data):
        del links[code]
        save_links(links)
        return jsonify({
            'success': False,
            'message': f'代码为 {code} 的短链接已过期'
        }), 410
    
    # 增加访问次数
    link_data['visit_count'] = link_data.get('visit_count', 0) + 1
    links[code] = link_data
    save_links(links)
    
    target_url = link_data['target_url']
    
    # 确保URL包含协议
    if not target_url.startswith(('http://', 'https://')):
        target_url = 'http://' + target_url
    
    return jsonify({
        'success': True,
        'code': code,
        'target_url': target_url
    })

@app.route('/api/stats/<code>', methods=['GET'])
def get_link_stats(code):
    """获取链接统计信息"""
    if len(code) != 4:
        return jsonify({
            'success': False,
            'message': '链接代码必须是4位字符'
        }), 404
    
    links = load_links()
    if code not in links:
        return jsonify({
            'success': False,
            'message': f'未找到代码为 {code} 的短链接'
        }), 404
    
    link_data = links[code]
    
    # 检查是否过期
    if is_link_expired(link_data):
        return jsonify({
            'success': False,
            'message': f'代码为 {code} 的短链接已过期',
            'expired': True
        }), 410
    
    return jsonify({
        'success': True,
        'code': code,
        'target_url': link_data['target_url'],
        'visit_count': link_data.get('visit_count', 0),
        'created_at': link_data.get('created_at'),
        'expires_at': link_data.get('expires_at'),
        'expire_option': link_data.get('expire_option')
    })

@app.route('/api/admin/delete/<code>', methods=['DELETE'])
def admin_delete_link(code):
    """管理员删除指定链接"""
    if len(code) != 4:
        return jsonify({
            'success': False,
            'message': '链接代码必须是4位字符'
        }), 400
    
    links = load_links()
    if code not in links:
        return jsonify({
            'success': False,
            'message': f'未找到代码为 {code} 的短链接'
        }), 404
    
    # 删除链接
    del links[code]
    save_links(links)
    
    return jsonify({
        'success': True,
        'message': f'成功删除短链接 {code}'
    })

@app.route('/api/admin/delete-batch', methods=['POST'])
def admin_delete_batch():
    """管理员批量删除链接"""
    data = request.json
    codes = data.get('codes', [])
    
    if not codes or not isinstance(codes, list):
        return jsonify({
            'success': False,
            'message': '请提供要删除的链接代码列表'
        }), 400
    
    links = load_links()
    deleted = []
    not_found = []
    
    for code in codes:
        if len(code) == 4 and code in links:
            del links[code]
            deleted.append(code)
        else:
            not_found.append(code)
    
    save_links(links)
    
    return jsonify({
        'success': True,
        'deleted': deleted,
        'deleted_count': len(deleted),
        'not_found': not_found,
        'message': f'成功删除 {len(deleted)} 个短链接'
    })

@app.route('/api/admin/clear-expired', methods=['POST'])
def admin_clear_expired():
    """管理员清理所有过期链接"""
    expired_count = clean_expired_links()
    
    return jsonify({
        'success': True,
        'cleared_count': expired_count,
        'message': f'成功清理 {expired_count} 个过期链接'
    })

#==========================后端API接口==========================#


if __name__ == '__main__':
    # 确保数据文件存在
    if not os.path.exists(DATA_FILE):
        save_links({})
    
    app.run(debug=True, host='0.0.0.0', port=5000)