app/services.py

import os
import logging
import random
import zipfile
import tempfile
import shutil
import json
from typing import List, Dict, Any, Optional, Union
from PIL import Image
from app.api import get_chat_completion
from app.config import (
    STICKER_RERANKING_SYSTEM_PROMPT,
    PUBLIC_URL,
    TEMP_DIR
)

from app.database import db
from app.image_utils import (
    save_image_temp, 
    generate_temp_image,
    upload_folder_to_huggingface,
    upload_to_huggingface, 
    get_image_cdn_url, 
    get_image_description,
    calculate_image_hash
)
from app.gradio_formatter import gradio_formatter
from multiprocessing import Queue

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
logger = logging.getLogger(__name__)


class StickerService:
    """贴纸服务类，处理贴纸的上传、搜索等业务逻辑"""
    
    @staticmethod
    def upload_sticker(image_file_path: str, title: str, description: str, tags: str) -> str:
        """上传贴纸"""
        try:
            # 打开图片
            image = Image.open(image_file_path)

            # 检查文件名是否已存在
            image_hash = calculate_image_hash(image)
            if db.check_image_exists(image_hash):
                print(f"文件已存在", image_hash)
                raise Exception('File_Exists')
            
            # 上传到 HuggingFace
            file_path, image_filename = upload_to_huggingface(image_file_path)
            # print('>>>> image_file_path', image_file_path)
            # print('>>>> image_filename', image_filename)
            # print('>>>> file_path', file_path)

            # 如果没有描述，获取图片描述
            if not description:
                image_cdn_url = ''
                if (PUBLIC_URL):
                    image_cdn_url = f'{PUBLIC_URL}/gradio_api/file={image_file_path}'
                else:
                    image_cdn_url = get_image_cdn_url(file_path)
                print('image_cdn_url',image_cdn_url)
                
                description = get_image_description(image_cdn_url)
            
            # 清理临时文件
            # os.unlink(temp_file_path)
            
            # 存储到 Milvus
            db.store_sticker(title, description, tags, file_path, image_hash)
            
            return f"Upload successful! {image_filename}"
            
        except Exception as e:
            logger.error(f"Upload failed: {str(e)}")
            return f"Upload failed: {str(e)}"
    
    @staticmethod
    def import_stickers(
        sticker_dataset: str,
        upload: bool = False,
        save_to_milvus: bool = False,
        progress_callback: callable = None,
    ) -> List[str]:
        """导入表情包数据集
        
        Args:
            sticker_dataset (str): 表情包数据集路径
            upload (bool, optional): 是否上传到HuggingFace. Defaults to False.
            progress_callback (callable, optional): 进度回调函数. Defaults to None.
        """
        results = []
        descriptions = {}
        
        try:
            # 创建临时目录
            cache_folder = os.path.join(TEMP_DIR, 'cache/')
            img_folder = os.path.join(TEMP_DIR, 'images/')
            data_json_path = os.path.join(cache_folder, 'data.json')
            stickers = []
            logger.info(f"start import dataset")
            # 解压数据集
            with zipfile.ZipFile(sticker_dataset, 'r') as zip_ref:
                zip_ref.extractall(cache_folder)
            logger.info(f"Extracted dataset to: {cache_folder}")
            
            # 尝试读取data.json文件
            if os.path.exists(data_json_path):
                with open(data_json_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    descriptions = { x["filename"]: x["content"] for x in data }
                logger.info(f"Loaded descriptions from data.json")
            
            # 遍历解压后的目录
            for root, dirs, files in os.walk(cache_folder):
                for file in files:
                    if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp')):
                        image_path = os.path.join(root, file)
                        
                        try:
                            # 打开图片
                            image = Image.open(image_path)
                            image_hash = calculate_image_hash(image)
                            if db.check_image_exists(image_hash):
                                results.append(f"跳过已存在的图片: {file}")
                                if progress_callback:
                                    progress_callback(file, "Skipped (exists)")
                                continue
                            
                            # 获取图片描述
                            description = None
                            if file in descriptions:
                                description = descriptions[file]

                            if not description:
                                results.append(f"跳过无描述的图片: {file}")
                                if progress_callback:
                                    progress_callback(file, "Skipped (no description)")
                                continue
                            
                            image_filename = f"image_{random.randint(100000, 999999)}.png"
                            file_path = f"images/{image_filename}"

                            generate_temp_image(img_folder, image, image_filename)
                            if save_to_milvus:
                                db.store_sticker("", description, "", file_path, image_hash)
                            stickers.append({
                                "title": "",
                                "description": description,
                                "tags": "",
                                "file_path": file_path,
                                "image_hash": image_hash
                            })
                            if progress_callback:
                                results.append(f"成功导入: {image_filename}")
                                progress_callback(file, "Imported")
                            
                        except Exception as e:
                            logger.error(f"Failed to process image {file}: {str(e)}")
                            results.append(f"处理失败 {file}: {str(e)}")
                            if progress_callback:
                                progress_callback(file, f"Failed: {str(e)}")

            # 上传到 HuggingFace
            if upload and len(stickers) > 0:
                logger.info(f"upload to huggingface, {len(stickers)} stickers")
                upload_folder_to_huggingface(img_folder)
                results.append(f"上传到 HuggingFace 成功")

            return results
            
        except Exception as e:
            logger.error(f"Import failed: {str(e)}")
            results.append(f"导入失败: {str(e)}")
            return results
            
        finally:
            # 清理临时目录
            if cache_folder and os.path.exists(cache_folder):
                shutil.rmtree(cache_folder)
                logger.info(f"Cleaned up temporary directory: {cache_folder}")

            if img_folder and os.path.exists(img_folder):
                shutil.rmtree(img_folder)
                logger.info(f"Cleaned up temporary directory: {img_folder}")

    @staticmethod
    def search_stickers(description: str, limit: int = 2, reranking : bool = False) -> List[Dict[str, Any]]:
        """搜索贴纸"""
        if not description:
            return []
        
        try:
            results = db.search_stickers(description, limit)
            if (reranking):
                # 对搜索结果进行重排
                results = StickerService.rerank_search_results(description, results, limit)
            return results
        except Exception as e:
            logger.error(f"Search failed: {str(e)}")
            return []
    
    @staticmethod
    def get_all_stickers(limit: int = 1000) -> List[List]:
        """获取所有贴纸"""
        try:
            results = db.get_all_stickers(limit)
            return gradio_formatter.format_all_stickers(results)
        except Exception as e:
            logger.error(f"Failed to get all stickers: {str(e)}")
            return []
            
    @staticmethod
    def delete_sticker(sticker_id: str) -> str:
        """删除贴纸"""
        try:
            # 首先查询贴纸是否存在
            result = db.delete_sticker(sticker_id)
            
            return f"Sticker with ID {sticker_id} deleted successfully"
        except Exception as e:
            logger.error(f"Delete failed: {str(e)}")
            return f"Delete failed: {str(e)}"
    

    @staticmethod
    def rerank_search_results(query: str, sticker_list: List[Dict[str, Any]], limit: int = 5) -> List[Dict[str, Any]]:
        ## 使用 LLM 模型重新排序搜索结果
        try:
            # 构建提示词
            system_prompt = STICKER_RERANKING_SYSTEM_PROMPT
            # 构建用户提示词，包含查询和表情包信息
            _sticker_list = []
            for hit in sticker_list:
                _sticker_list.append({
                    "id": hit["id"],
                    "description": hit["entity"]["description"]
                })
            
            user_prompt = f"请分析关键词 '{query}' 与以下表情包的相关性:\n{_sticker_list}"
            
            print(f">>> 使用 LLM 模型重新排序....", user_prompt, system_prompt)
            # 调用 LLM 模型获取重排序结果
            response = get_chat_completion(user_prompt, system_prompt)
            
            # 解析 LLM 返回的 JSON 结果
            reranked_stickers = json.loads(response)
            
            # 验证返回结果格式
            if not isinstance(reranked_stickers, list):
                raise ValueError("Invalid response format")
                
            # 按分数排序
            reranked_stickers.sort(key=lambda x: float(x.get("score", 0)), reverse=True)

            print(f">>> LLM 排序结果", reranked_stickers)

            
            # 将重排序结果与原始结果对应
            rerank_results = []
            for sticker in reranked_stickers:
                for hit in sticker_list:
                    if str(hit["id"]) == str(sticker["sticker_id"]):
                        hit["entity"]["score"] = sticker["score"]
                        hit["entity"]["reason"] = sticker["reason"]
                        rerank_results.append(hit)
                        break
            
            print(f">>> rerank_results", rerank_results)
            return rerank_results

        except Exception as e:
            logger.error(f"Reranking failed: {str(e)}")
            return []




# 创建服务实例
sticker_service = StickerService()