#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
生产级Web应用 - 使用纯百度VDB系统，完全移除FAISS依赖
优化版本：支持自动清理、内存处理和流式上传
"""

import os
import sys
import json
import base64
import time
import logging
from io import BytesIO
from pathlib import Path
from typing import Dict, List, Any, Optional

import numpy as np
from PIL import Image
from flask import Flask, request, jsonify, render_template, send_from_directory
from werkzeug.utils import secure_filename
import threading

from multimodal_retrieval_vdb_only import MultimodalRetrievalVDBOnly
from mongodb_manager import get_mongodb_manager
from baidu_bos_manager import get_bos_manager
from optimized_file_handler import get_file_handler

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 创建Flask应用
app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max file size

# 全局变量
retrieval_system = None
system_lock = threading.Lock()

# 配置
UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'}

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

def allowed_file(filename):
    """检查文件扩展名是否允许"""
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def encode_image_to_base64(image_path: str) -> str:
    """将图像编码为base64字符串"""
    try:
        with open(image_path, 'rb') as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
            return f"data:image/jpeg;base64,{encoded_string}"
    except Exception as e:
        logger.error(f"图像编码失败: {e}")
        return ""

@app.route('/')
def index():
    """主页"""
    return render_template('index.html')

@app.route('/api/init', methods=['POST'])
def init_system():
    """初始化多模态检索系统"""
    global retrieval_system
    
    try:
        with system_lock:
            if retrieval_system is None:
                logger.info("🚀 初始化纯VDB多模态检索系统...")
                retrieval_system = MultimodalRetrievalVDBOnly()
                logger.info("✅ 系统初始化完成")
            
            return jsonify({
                'success': True,
                'message': '系统初始化成功',
                'backend': 'Baidu VDB (No FAISS)',
                'status': 'ready'
            })
    
    except Exception as e:
        logger.error(f"❌ 系统初始化失败: {e}")
        return jsonify({
            'success': False,
            'message': f'系统初始化失败: {str(e)}'
        }), 500

@app.route('/api/upload/texts', methods=['POST'])
@app.route('/api/upload_texts', methods=['POST'])
def upload_texts():
    """上传文本数据 - 优化版本"""
    global retrieval_system
    
    try:
        if retrieval_system is None:
            return jsonify({
                'success': False,
                'message': '系统未初始化，请先调用 /api/init'
            }), 400
        
        data = request.get_json()
        texts = data.get('texts', [])
        
        if not texts:
            return jsonify({
                'success': False,
                'message': '没有提供文本数据'
            }), 400
        
        # 获取优化的文件处理器
        file_handler = get_file_handler()
        
        # 使用内存处理文本数据
        logger.info(f"📝 开始内存处理 {len(texts)} 条文本...")
        processed_texts = file_handler.process_text_in_memory(texts)
        
        if not processed_texts:
            return jsonify({
                'success': False,
                'message': '文本处理失败'
            }), 400
        
        # 自动构建文本索引
        index_status = "索引构建失败"
        try:
            with system_lock:
                retrieval_system.build_text_index_parallel(texts)
                
                # 存储向量元数据
                mongodb_mgr = get_mongodb_manager()
                for text_info in processed_texts:
                    mongodb_mgr.store_vector_metadata(
                        file_id=text_info["file_id"],
                        vector_type="text_vector",
                        vdb_id=text_info["bos_key"],
                        vector_info={"text_content": text_info["text_content"]}
                    )
                    
            logger.info(f"✅ 文本索引自动构建完成")
            index_status = "索引已自动构建"
        except Exception as e:
            logger.error(f"❌ 自动构建文本索引失败: {e}")
            index_status = f"索引构建失败: {str(e)}"
        
        return jsonify({
            'success': True,
            'message': f'成功处理 {len(texts)} 条文本（内存模式）',
            'count': len(texts),
            'processed_texts': len(processed_texts),
            'index_status': index_status,
            'auto_indexed': True,
            'processing_method': 'memory',
            'storage_info': {
                'bos_stored': len(processed_texts),
                'mongodb_stored': len(processed_texts)
            }
        })
    
    except Exception as e:
        logger.error(f"❌ 文本上传失败: {e}")
        return jsonify({
            'success': False,
            'message': f'文本上传失败: {str(e)}'
        }), 500

@app.route('/api/upload/images', methods=['POST'])
@app.route('/api/upload_images', methods=['POST'])
def upload_images():
    """上传图像文件 - 优化版本"""
    global retrieval_system
    
    try:
        if retrieval_system is None:
            return jsonify({
                'success': False,
                'message': '系统未初始化，请先调用 /api/init'
            }), 400
        
        if 'files' not in request.files:
            return jsonify({
                'success': False,
                'message': '没有文件上传'
            }), 400
        
        files = request.files.getlist('files')
        if not files or all(file.filename == '' for file in files):
            return jsonify({
                'success': False,
                'message': '没有选择文件'
            }), 400
        
        # 获取优化的文件处理器
        file_handler = get_file_handler()
        processed_files = []
        temp_files_for_model = []
        
        # 处理每个文件
        for file in files:
            if file and allowed_file(file.filename):
                filename = secure_filename(file.filename)
                logger.info(f"📷 处理图像文件: {filename}")
                
                try:
                    # 智能处理图像（自动选择内存或临时文件）
                    result = file_handler.process_image_smart(file, filename)
                    
                    if result:
                        processed_files.append(result)
                        
                        # 如果需要为模型处理准备临时文件
                        if result.get('processing_method') == 'memory':
                            # 对于内存处理的文件，需要为模型创建临时文件
                            temp_path = file_handler.get_temp_file_for_model(file, filename)
                            if temp_path:
                                temp_files_for_model.append({
                                    'temp_path': temp_path,
                                    'file_info': result
                                })
                        else:
                            # 临时文件处理的情况，直接使用返回的路径
                            temp_files_for_model.append({
                                'temp_path': result.get('temp_path'),
                                'file_info': result
                            })
                        
                        logger.info(f"✅ 图像处理成功: {filename} ({result['processing_method']})")
                    else:
                        logger.error(f"❌ 图像处理失败: {filename}")
                        
                except Exception as e:
                    logger.error(f"❌ 处理图像文件失败 {filename}: {e}")
        
        if processed_files:
            # 自动构建图像索引
            index_status = "索引构建失败"
            try:
                # 准备图像路径列表
                image_paths = [item['temp_path'] for item in temp_files_for_model if item['temp_path']]
                
                if image_paths:
                    with system_lock:
                        retrieval_system.build_image_index_parallel(image_paths)
                        
                        # 存储向量元数据
                        mongodb_mgr = get_mongodb_manager()
                        for file_info in processed_files:
                            mongodb_mgr.store_vector_metadata(
                                file_id=file_info["file_id"],
                                vector_type="image_vector",
                                vdb_id=file_info["bos_key"],
                                vector_info={"filename": file_info["filename"]}
                            )
                    
                    logger.info(f"✅ 图像索引自动构建完成")
                    index_status = "索引已自动构建"
                    
                    # 清理模型处理用的临时文件
                    for item in temp_files_for_model:
                        if item['temp_path']:
                            file_handler.cleanup_temp_file(item['temp_path'])
                    
            except Exception as e:
                logger.error(f"❌ 自动构建图像索引失败: {e}")
                index_status = f"索引构建失败: {str(e)}"
                
                # 即使索引失败也要清理临时文件
                for item in temp_files_for_model:
                    if item['temp_path']:
                        file_handler.cleanup_temp_file(item['temp_path'])
            
            return jsonify({
                'success': True,
                'message': f'成功上传 {len(processed_files)} 个图像文件',
                'count': len(processed_files),
                'processed_files': len(processed_files),
                'index_status': index_status,
                'auto_indexed': True,
                'processing_methods': [f['processing_method'] for f in processed_files],
                'storage_info': {
                    'bos_stored': len(processed_files),
                    'mongodb_stored': len(processed_files)
                }
            })
        else:
            return jsonify({
                'success': False,
                'message': '没有有效的图像文件'
            }), 400
    
    except Exception as e:
        logger.error(f"❌ 图像上传失败: {e}")
        # 确保清理所有临时文件
        file_handler = get_file_handler()
        file_handler.cleanup_all_temp_files()
        return jsonify({
            'success': False,
            'message': f'图像上传失败: {str(e)}'
        }), 500

@app.route('/api/search/image_to_text', methods=['POST'])
def search_image_to_text():
    """图搜文 - 优化版本"""
    global retrieval_system
    
    try:
        if retrieval_system is None:
            return jsonify({
                'success': False,
                'message': '系统未初始化'
            }), 400
        
        top_k = request.form.get('top_k', 5, type=int)
        
        if 'image' not in request.files:
            return jsonify({
                'success': False,
                'message': '没有上传图像'
            }), 400
        
        file = request.files['image']
        if not file or not allowed_file(file.filename):
            return jsonify({
                'success': False,
                'message': '无效的图像文件'
            }), 400
        
        # 获取优化的文件处理器
        file_handler = get_file_handler()
        
        # 为模型处理创建临时文件
        filename = secure_filename(file.filename)
        temp_path = file_handler.get_temp_file_for_model(file, filename)
        
        if not temp_path:
            return jsonify({
                'success': False,
                'message': '临时文件创建失败'
            }), 500
        
        try:
            # 执行图搜文
            with system_lock:
                results = retrieval_system.search_image_to_text(temp_path, top_k=top_k)
            
            # 格式化结果
            formatted_results = []
            for text, score in results:
                formatted_results.append({
                    'text': text,
                    'score': float(score),
                    'type': 'text'
                })
            
            # 编码查询图像
            query_image_base64 = encode_image_to_base64(temp_path)
            
            return jsonify({
                'success': True,
                'query_image': query_image_base64,
                'results': formatted_results,
                'count': len(formatted_results),
                'search_type': 'image_to_text',
                'processing_method': 'optimized_temp_file'
            })
            
        finally:
            # 确保清理临时文件
            file_handler.cleanup_temp_file(temp_path)
    
    except Exception as e:
        logger.error(f"❌ 图搜文失败: {e}")
        return jsonify({
            'success': False,
            'message': f'图搜文失败: {str(e)}'
        }), 500

@app.route('/api/search/image_to_image', methods=['POST'])
def search_image_to_image():
    """图搜图 - 优化版本"""
    global retrieval_system
    
    try:
        if retrieval_system is None:
            return jsonify({
                'success': False,
                'message': '系统未初始化'
            }), 400
        
        top_k = request.form.get('top_k', 5, type=int)
        
        if 'image' not in request.files:
            return jsonify({
                'success': False,
                'message': '没有上传图像'
            }), 400
        
        file = request.files['image']
        if not file or not allowed_file(file.filename):
            return jsonify({
                'success': False,
                'message': '无效的图像文件'
            }), 400
        
        # 获取优化的文件处理器
        file_handler = get_file_handler()
        
        # 为模型处理创建临时文件
        filename = secure_filename(file.filename)
        temp_path = file_handler.get_temp_file_for_model(file, filename)
        
        if not temp_path:
            return jsonify({
                'success': False,
                'message': '临时文件创建失败'
            }), 500
        
        try:
            # 执行图搜图
            with system_lock:
                results = retrieval_system.search_image_to_image(temp_path, top_k=top_k)
            
            # 格式化结果
            formatted_results = []
            for image_path, score in results:
                image_base64 = encode_image_to_base64(image_path)
                formatted_results.append({
                    'image_path': image_path,
                    'image_base64': image_base64,
                    'score': float(score),
                    'type': 'image'
                })
            
            # 编码查询图像
            query_image_base64 = encode_image_to_base64(temp_path)
            
            return jsonify({
                'success': True,
                'query_image': query_image_base64,
                'results': formatted_results,
                'count': len(formatted_results),
                'search_type': 'image_to_image',
                'processing_method': 'optimized_temp_file'
            })
            
        finally:
            # 确保清理临时文件
            file_handler.cleanup_temp_file(temp_path)
    
    except Exception as e:
        logger.error(f"❌ 图搜图失败: {e}")
        return jsonify({
            'success': False,
            'message': f'图搜图失败: {str(e)}'
        }), 500

@app.route('/api/data/list', methods=['GET'])
def list_data():
    """获取数据列表"""
    global retrieval_system
    
    try:
        if retrieval_system is None:
            return jsonify({
                'success': False,
                'message': '系统未初始化'
            }), 400
        
        # 获取MongoDB管理器
        mongodb_mgr = get_mongodb_manager()
        
        # 获取所有文件元数据
        files_data = mongodb_mgr.get_all_files()
        
        images = []
        texts = []
        
        for file_data in files_data:
            if file_data.get('file_type') == 'image':
                # 提取文件名用于显示
                file_path = file_data.get('file_path', '')
                filename = os.path.basename(file_path) if file_path else file_data.get('bos_key', '')
                images.append(filename)
            elif file_data.get('file_type') == 'text':
                # 获取文本内容
                text_content = file_data.get('additional_info', {}).get('text_content', '')
                if text_content:
                    texts.append(text_content)
        
        return jsonify({
            'success': True,
            'images': images,
            'texts': texts,
            'total_files': len(files_data),
            'image_count': len(images),
            'text_count': len(texts)
        })
    
    except Exception as e:
        logger.error(f"❌ 获取数据列表失败: {e}")
        return jsonify({
            'success': False,
            'message': f'获取数据列表失败: {str(e)}'
        }), 500

@app.route('/api/stats', methods=['GET'])
@app.route('/api/status', methods=['GET'])
@app.route('/api/data/stats', methods=['GET'])
def get_stats():
    """获取系统统计信息"""
    global retrieval_system
    
    try:
        if retrieval_system is None:
            return jsonify({
                'success': False,
                'message': '系统未初始化'
            }), 400
        
        # 获取MongoDB统计信息
        mongodb_mgr = get_mongodb_manager()
        mongodb_stats = mongodb_mgr.get_stats()
        
        with system_lock:
            stats = retrieval_system.get_statistics()
            gpu_info = retrieval_system.get_gpu_memory_info()
        
        # 合并统计信息
        enhanced_stats = {
            **stats,
            'mongodb_stats': mongodb_stats,
            'storage_backend': {
                'metadata': 'MongoDB',
                'files': 'Baidu BOS'
            }
        }
        
        return jsonify({
            'success': True,
            'stats': enhanced_stats,
            'gpu_info': gpu_info,
            'backend': 'Baidu VDB + MongoDB + BOS'
        })
    
    except Exception as e:
        logger.error(f"❌ 获取统计信息失败: {e}")
        return jsonify({
            'success': False,
            'message': f'获取统计信息失败: {str(e)}'
        }), 500

@app.route('/api/clear', methods=['POST'])
def clear_data():
    """清空所有数据 - 优化版本"""
    global retrieval_system
    
    try:
        if retrieval_system is None:
            return jsonify({
                'success': False,
                'message': '系统未初始化'
            }), 400
        
        with system_lock:
            retrieval_system.clear_all_data()
        
        # 清理所有临时文件
        file_handler = get_file_handler()
        file_handler.cleanup_all_temp_files()
        
        return jsonify({
            'success': True,
            'message': '所有数据已清空，临时文件已清理'
        })
    
    except Exception as e:
        logger.error(f"❌ 清空数据失败: {e}")
        return jsonify({
            'success': False,
            'message': f'清空数据失败: {str(e)}'
        }), 500

@app.route('/uploads/<filename>')
def uploaded_file(filename):
    """提供上传文件的访问"""
    return send_from_directory(UPLOAD_FOLDER, filename)

@app.errorhandler(413)
def too_large(e):
    """文件过大错误处理"""
    return jsonify({
        'success': False,
        'message': '文件过大，最大支持16MB'
    }), 413

@app.errorhandler(500)
def internal_error(e):
    """内部服务器错误处理"""
    return jsonify({
        'success': False,
        'message': '内部服务器错误'
    }), 500

def init_app():
    """应用初始化 - 优化版本"""
    global retrieval_system
    
    print("=" * 60)
    print("生产级多模态检索Web应用 (优化版)")
    print("Backend: 纯百度VDB (无FAISS)")
    print("Features: 自动清理 + 内存处理 + 流式上传")
    print("=" * 60)
    
    # 显示GPU信息
    if torch.cuda.is_available():
        gpu_count = torch.cuda.device_count()
        print(f"检测到 {gpu_count} 个GPU:")
        for i in range(gpu_count):
            gpu_name = torch.cuda.get_device_properties(i).name
            gpu_memory = torch.cuda.get_device_properties(i).total_memory / (1024**3)
            print(f"  GPU {i}: {gpu_name} ({gpu_memory:.1f}GB)")
    
    print("=" * 60)
    
    # 启动时自动初始化系统
    try:
        logger.info("🚀 启动时自动初始化优化版VDB检索系统...")
        retrieval_system = MultimodalRetrievalVDBOnly()
        logger.info("✅ 优化版系统自动初始化成功")
        
        # 初始化文件处理器
        file_handler = get_file_handler()
        logger.info("✅ 优化文件处理器初始化成功")
        
    except Exception as e:
        logger.error(f"❌ 系统自动初始化失败: {e}")
        logger.info("系统将在首次API调用时初始化")

if __name__ == '__main__':
    import torch
    
    # 初始化应用
    init_app()
    
    # 启动Flask应用
    app.run(host='0.0.0.0', port=5000, debug=False, threaded=True)