from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from pymongo import MongoClient
from urllib.parse import quote_plus
import uuid
from typing import List, Optional
import json
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import HTMLResponse
import os
import base64
from groq import Groq
import faiss
import pickle
import torch
from transformers import CLIPProcessor, CLIPModel
from PIL import Image

# Load the FAISS index
index = faiss.read_index("knowledge_base.faiss")

# Load the titles metadata
with open("titles.pkl", "rb") as f:
    titles = pickle.load(f)

# Load CLIP model and processor on CPU
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to("cpu")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# Initialize Groq client
client = Groq(api_key='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')

# MongoDB connection setup
def get_mongo_client():
    password = quote_plus("momimaad@123")  # Change this to your MongoDB password
    mongo_uri = f"mongodb+srv://hammad:{password}@cluster0.2a9yu.mongodb.net/"
    return MongoClient(mongo_uri)

db_client = get_mongo_client()
db = db_client["recipe"]
user_collection = db["user_info"]

# Pydantic models for user data
class User(BaseModel):
    first_name: str
    last_name: str
    email: str
    password: str

class UserData(BaseModel):
    email: str
    password: str

class UserToken(BaseModel):
    token: str

class RecipeData(BaseModel):
    name: str

class AltrecipeData(BaseModel):
    recipe_name: str
    dietary_restrictions: str
    allergies: List

class Ingredient(BaseModel):
    name: str
    quantity: str


class Recipe(BaseModel):
    recipe_name: str
    ingredients: List[Ingredient]
    directions: List[str]


class get_recipe_name(BaseModel):
    recipe_name: List[str]
    ingredients: List[List[str]]

# Data model for LLM to generate
class Alternative_Ingredient(BaseModel):
    name: str
    quantity: str

class Alternative_Recipe(BaseModel):
    recipe_name: str
    alternative_ingredients: List[Alternative_Ingredient]
    alternative_directions: List[str]

# Function for finding the most similar image
def find_similar_image(image_path, threshold=30.0):
    # Load and preprocess the input image
    image = Image.open(image_path).convert("RGB")
    inputs = processor(images=image, return_tensors="pt")  

    # Generate embedding for the input image on CPU
    with torch.no_grad():
        image_features = model.get_image_features(**inputs).numpy()  # No need for .cpu()

    # Perform similarity search in FAISS
    distances, indices = index.search(image_features, k=1)  # Search for the most similar embedding

    # Check if the closest match meets the threshold
    if distances[0][0] < threshold:
        return titles[indices[0][0]]
    else:
        return "Not Found"
def get_recipe(recipe_name: str) -> Recipe:
    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "system",
                "content": f"""Your are an expert agent to generate a recipes with proper and corrected ingredients and direction. Your directions should be concise and to the point and dont explain any irrelevant text.
                You are a recipe database that outputs recipes in JSON.\n
              The JSON object must use the schema: {json.dumps(Recipe.model_json_schema(), indent=2)}""",
            },
            {
                "role": "user",
                "content": f"Fetch a recipe for {recipe_name}",
            },
        ],
        model="llama-3.2-90b-text-preview",
        temperature=0,
        # Streaming is not supported in JSON mode
        stream=False,
        # Enable JSON mode by setting the response format
        response_format={"type": "json_object"},
    )
    return Recipe.model_validate_json(chat_completion.choices[0].message.content)


def Suggest_ingredient_alternatives(recipe_name: str, dietary_restrictions: str, allergies: List) -> Alternative_Recipe:
    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "system",
                "content": f"""
                 You are an expert agent to suggest alternatives for specific allergies ingredients for the provided recipe {recipe_name}.

                Please take the following into account:
                - If the user has dietary restrictions, suggest substitutes that align with their needs (e.g., vegan, gluten-free, etc.) in alternative_directions and your alternative_directions should be concise and to the point.
                -In ingredient you will recommend the safe ingredient for avoid any allergy and dietary restriction.
                - Consider the following allergies {allergies} and recommend the safe ingredient to avoid this allergies.

                recipe_name: {recipe_name}
                Dietary Restrictions: {dietary_restrictions}
                Allergies: {', '.join(allergies)}

                You are a recipe database that outputs alternative recipes to avoid allergy and dietary_restrictions in JSON.\n
                The JSON object must use the schema: {json.dumps(Alternative_Recipe.model_json_schema(), indent=2)}""",
            },
            {
                "role": "user",
                "content": f"""Fetch a alternative recipe for recipe_name: {recipe_name}
                Dietary Restrictions: {dietary_restrictions}
                Allergies: {', '.join(allergies)}""",
            },
        ],
        model="llama-3.2-90b-text-preview",
        temperature=0,
        # Streaming is not supported in JSON mode
        stream=False,
        # Enable JSON mode by setting the response format
        response_format={"type": "json_object"},
    )
    return Alternative_Recipe.model_validate_json(chat_completion.choices[0].message.content)

app = FastAPI()

@app.post("/get_recipe/{token}")
async def get_recipe_response(token: str, recipe_user: RecipeData):
    user = user_collection.find_one({"token": token})
    if not user:
        raise HTTPException(status_code=401, detail="Invalid token")

    # Find user by email
    recipe_name = recipe_user.name
    response = get_recipe(recipe_name)
    return {
        "Response": response
    }

@app.post("/get_recipe_alternative/{token}")
async def get_alternative_recipe_response(token: str, altrecipe_user: AltrecipeData):
    user = user_collection.find_one({"token": token})
    if not user:
        raise HTTPException(status_code=401, detail="Invalid token")

    response = Suggest_ingredient_alternatives(altrecipe_user.recipe_name, altrecipe_user.dietary_restrictions, altrecipe_user.allergies)
    return {
        "Response": response
    }


# Directory to save uploaded images
UPLOAD_DIR = "uploads"

# Ensure the upload directory exists
os.makedirs(UPLOAD_DIR, exist_ok=True)


# Endpoint to upload an image
@app.post("/upload-image/{token}")
async def upload_image(token: str, file: UploadFile = File(...)):
    user = user_collection.find_one({"token": token})
    if not user:
        raise HTTPException(status_code=401, detail="Invalid token")

    # Validate the file type
    if not file.filename.lower().endswith(('.png', '.jpg', '.jpeg')):
        raise HTTPException(status_code=400, detail="Invalid file type. Only PNG, JPG, and JPEG are allowed.")

    # Create a file path for saving the uploaded file
    file_path = os.path.join(UPLOAD_DIR, file.filename)

    # Save the file
    with open(file_path, "wb") as buffer:
        buffer.write(await file.read())

    result = find_similar_image(file_path, threshold=30.0)

    return {
        "Response": result
    }


# Endpoint to register a new user
@app.post("/register")
async def register_user(user: User):
    # Check if user already exists
    existing_user = user_collection.find_one({"email": user.email})
    if existing_user:
        raise HTTPException(status_code=400, detail="Email already registered")

    # Create user data
    user_data = {
        "first_name": user.first_name,
        "last_name": user.last_name,
        "email": user.email,
        "password": user.password,  # Store plaintext password (not recommended in production)
    }

    # Insert the user data into the user_info collection
    result = user_collection.insert_one(user_data)
    return {"msg": "User registered successfully", "user_id": str(result.inserted_id)}

# Endpoint to check user credentials and generate a token
@app.post("/get_token")
async def check_credentials(user: UserData):
    # Find user by email
    existing_user = user_collection.find_one({"email": user.email})

    # Check if user exists and password matches
    if not existing_user or existing_user["password"] != user.password:
        raise HTTPException(status_code=401, detail="Invalid email or password")

    # Generate a UUID token
    token = str(uuid.uuid4())

    # Update the user document with the token
    user_collection.update_one({"email": user.email}, {"$set": {"token": token}})

    return {
        "first_name": existing_user["first_name"],
        "last_name": existing_user["last_name"],
        "token": token,
    }


@app.get("/")
async def root():
    return {"message": "API is up and running!"}