import os
from flask import Flask, request, jsonify, render_template
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
from openai import OpenAI
from flask_cors import CORS
from fridgeItemDetector import fridge_item_detector, black_list_classes

app = Flask(__name__, static_url_path='/static')
CORS(app)  # 允许所有来源的跨域请求


# 初始化 Flask 应用
app = Flask(__name__)

# # 加载预训练的 ResNet50 模型
# model = models.resnet50(pretrained=True)

# # 重新定义最后的全连接层，适应 67 个类别
# num_features = model.fc.in_features
# model.fc = nn.Linear(num_features, 67)

# # 加载模型权重，但忽略 fc 层
# state_dict = torch.load('resnet50_food_classifier.pth', map_location=torch.device('cpu'))
# # 删除 state_dict 中全连接层的权重
# del state_dict['fc.weight']
# del state_dict['fc.bias']
# # 加载剩余的权重
# model.load_state_dict(state_dict, strict=False)

# model.eval()

# 定义数据预处理
# transform = transforms.Compose([
#     transforms.Resize((224, 224)),
#     transforms.ToTensor(),
#     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
# ])

open_ai_api_key = os.getenv("OPEN_AI_KEY")

# 设置 LLM
client = OpenAI(api_key=open_ai_api_key)
def generate_recipe(ingredients_text):
    prompt = f"Ingredients: {ingredients_text}. List 3 possible recipes with these ingredients and provide the essential cooking steps only."
    message = [{"role": "system", "content": prompt}]
    
    try:
        completion = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=message,
            temperature=0.38,
            max_tokens=500
        )
        # 返回的菜谱内容中，每个步骤用换行符分隔
        return completion.choices[0].message.content
    except Exception as e:
        return f"Error generating recipe: {str(e)}"


# 路由：主页
@app.route('/')
def home():
    return render_template('index.html')  # 前端页面

# 路由：上传图片并识别
@app.route('/upload', methods=['POST'])
def upload_image():
    try:
        # 检查是否有文件上传
        if 'fridge-image' not in request.files:
            return jsonify({"error": "No file part"})
        
        file = request.files['fridge-image']
        if file.filename == '':
            return jsonify({"error": "No selected file"})
        
        # try:
        #     # 尝试打开并转换图片
        #     img = Image.open(file).convert('RGB')
        #     img = transform(img).unsqueeze(0)  # 增加 batch 维度
        # except Exception as e:
        #     return jsonify({"error": f"Error processing image: {str(e)}"})

        try:
            # 进行模型推理
            result = fridge_item_detector.detect_objects(
                file,
                prob_cutoff=0.9,
                black_list_classes=black_list_classes,
                annotate_image=False,
                return_unique_label=True
            )
            detected_ingredient = ", ".join(result)
        except Exception as e:
            return jsonify({"error": f"Error during model inference: {str(e)}"})

        # 根据食材生成菜谱
        try:
            # ingredients = [ "apple", "asparagus", "aubergine", "bacon", "banana", "basil", "beans",
            #                 "beef", "beetroot", "bell pepper", "bitter gourd", "blueberries", "broccoli",
            #                 "cabbage", "carrot", "cauliflower", "cheese", "chicken", "chillies",
            #                 "chocolate", "coriander", "corn", "courgettes", "cream", "cucumber", "dates",
            #                 "egg", "flour", "garlic", "ginger", "green beans", "green chilies", "ham",
            #                 "juice", "lemon", "lettuce", "lime", "mango", "meat", "mineral water",
            #                 "mushroom", "olive", "onion", "orange", "parsley", "peach", "peas", "peppers",
            #                 "potato", "pumpkin", "red grapes", "red onion", "salami", "sauce", "sausage",
            #                 "shallot", "shrimp", "spinach", "spring onion", "strawberry", "sugar",
            #                 "sweet potato", "swiss butter", "swiss jam", "swiss yoghurt", "tomato", "watermelon"]  # 所有的食材
            
            # detected_ingredient = ingredients[predicted_class]
            recipe = generate_recipe(detected_ingredient)
        except Exception as e:
            return jsonify({"error": f"Error generating recipe: {str(e)}"})

        return jsonify({"ingredient": detected_ingredient, "recipe": recipe})
    
    except Exception as e:
        return jsonify({"error": f"Unexpected error: {str(e)}"})

if __name__ == '__main__':
    app.run(debug=True)