from flask import Flask, request, jsonify
import os
from urllib.request import urlretrieve
import time
from image_search import ImageSearchEngine

app = Flask(__name__)

# 配置常量
UPLOAD_FOLDER = 'static/images'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp'}

# 搜索参数默认值
TOP_K = 5  # 默认返回的最大结果数
MIN_SCORE = 0.0  # 默认最小相似度分数
MAX_SCORE = 100.0  # 默认最大相似度分数

os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# 初始化图像搜索引擎
search_engine = ImageSearchEngine()

@app.route("/", methods=["POST"])
def invoke():
    """处理图片上传请求"""
    try:
        # 获取并验证参数
        data = request.get_json()
        if not data:
            return jsonify({'error': '请求必须包含JSON数据'}), 400
        
        print(f"Received data: {data}")
        image_url = data.get('url')
        
        time_base = str(time.time() * 1000)
        image_path = os.path.join(UPLOAD_FOLDER, time_base + os.path.basename(image_url))
        urlretrieve(image_url, image_path)
        
        start_time = time.time()
        vector = search_engine._process_image(image_path)
        os.remove(image_path)
        return jsonify({"vector_spend_time": time.time() - start_time,"vector": vector.tolist()}),200
            
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5001, debug=False)