import os
import json
import numpy as np
import faiss
from flask import Flask, request, jsonify, send_file
from transformers import LlamaTokenizer, LlamaModel
import torch
from tqdm import tqdm

# Load tokenizer and model (global for both embedding generation and API)
model_name = "openlm-research/open_llama_3b_v2"
tokenizer = LlamaTokenizer.from_pretrained(model_name)
model = LlamaModel.from_pretrained(model_name)
# Set the pad_token to eos_token
tokenizer.pad_token = tokenizer.eos_token

# model.eval()
device = "cuda" if torch.cuda.is_available() else "cpu"

model = model.to(device)

app = Flask(__name__)

file_path = "products.json"
# Function to generate embeddings
def generate_embeddings(batch_size=64, max_length=512):
    # Load product data
    with open(file_path, "r") as f:
        products = json.load(f)

    # Initialize storage for embeddings
    embeddings = []
    product_ids = []

    # Process products in batches
    num_batches = len(products) // batch_size + (len(products) % batch_size > 0)
    for i in tqdm(range(num_batches), desc="Generating embeddings"):
        batch = products[i * batch_size : (i + 1) * batch_size]
        descriptions = [product["description"] for product in batch]
        ids = [product["id"] for product in batch]
        
        # Tokenize and move inputs to the device
        inputs = tokenizer(
            descriptions, return_tensors="pt", padding=True, truncation=True, max_length=max_length
        ).to(device)
        
        # Generate embeddings
        with torch.no_grad():
            outputs = model(**inputs, output_hidden_states=True)
        batch_embeddings = outputs.hidden_states[-1].mean(dim=1).cpu().numpy()
        
        # Append results
        embeddings.extend(batch_embeddings)
        product_ids.extend(ids)

    # Convert embeddings to NumPy array
    embeddings = np.array(embeddings)
    product_ids = np.array(product_ids)

    # Save embeddings to FAISS index
    dimension = embeddings.shape[1]  # Embedding size
    index = faiss.IndexFlatL2(dimension)
    index.add(embeddings)
    faiss.write_index(index, "product_index.faiss")
    np.save("product_ids.npy", product_ids)

    print(f"Embeddings for {len(product_ids)} products generated successfully!")
    return f"Embeddings for {len(product_ids)} products generated successfully!"

@app.route("/", methods=["GET"])
def home():
    return "Flask app is running!"


# API Endpoint to Trigger Embedding Generation
@app.route("/generate_embeddings", methods=["POST"])
def trigger_embedding_generation():
    print("creating embeddings")
    generate_embeddings()
    return jsonify({"message": "Embeddings generated successfully!"})

# API Endpoint to Get Similar Products
@app.route("/similar", methods=["POST"])
def get_similar_products():
    print("finding right product")
    index = faiss.read_index("product_index.faiss")
    product_ids = np.load("product_ids.npy")
    with open(file_path, "r") as f:
        products = {p["id"]: p for p in json.load(f)}

    data = request.json
    product_id = data.get("product_id")
    product_index = np.where(product_ids == product_id)[0]
    if len(product_index) == 0:
        return jsonify({"error": "Product ID not found"}), 404

    query_embedding = index.reconstruct(int(product_index[0]))

    distances, indices = index.search(np.expand_dims(query_embedding, axis=0), k=5)
    similar_ids = product_ids[indices[0]].tolist()
    similar_products = [products[pid] for pid in similar_ids if pid != product_id]
    return jsonify({"similar_products": similar_products})

@app.route("/download/<filename>", methods=["GET"])
def download_file(filename):
    try:
        # Validate allowed files
        if filename not in ["product_index.faiss", "product_ids.npy"]:
            return {"error": "File not found or not allowed"}, 404
        
        # Serve the file for download
        return send_file(filename, as_attachment=True)
    except Exception as e:
        return {"error": str(e)}, 500

def get_similar_produdcts_direct(product_id):
    print("finding right product")
    index = faiss.read_index("product_index.faiss")
    product_ids = np.load("product_ids.npy")
    with open(file_path, "r") as f:
        products = {p["id"]: p for p in json.load(f)}
    # print(products[:6])
    print(len(products))
    print(product_id)
    product_index = np.where(product_ids == product_id)[0]
    if len(product_index) == 0:
        return jsonify({"error": "Product ID not found"}), 404

    query_embedding = index.reconstruct(int(product_index[0]))

    distances, indices = index.search(np.expand_dims(query_embedding, axis=0), k=5)
    similar_ids = product_ids[indices[0]].tolist()
    similar_products = [products[pid] for pid in similar_ids if pid != product_id]
    print("Ran fine")
    print(similar_products)
    # return jsonify({"similar_products": similar_products})


if __name__ == "__main__":
    # Check if embeddings already exist; if not, generate them
    if not os.path.exists("product_index.faiss") or not os.path.exists("product_ids.npy"):
        generate_embeddings()
    # get_similar_produdcts_direct(2)
    app.run(host="0.0.0.0", port=7860)