import os
import warnings
from transformers import AutoModelForImageClassification, AutoFeatureExtractor
import torch
from flask_cors import CORS
from flask import Flask, request, json, Response
import numpy as np
from PIL import Image
import requests
from io import BytesIO
from bs4 import BeautifulSoup
from urllib.parse import urljoin
import google.generativeai as genai
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_google_genai import ChatGoogleGenerativeAI
import warnings
from dotenv import load_dotenv

os.environ["CUDA_VISIBLE_DEVICES"] = ""

app = Flask(__name__)
cors = CORS(app)
load_dotenv()

# # Define the model and feature extractor globally
# model = AutoModelForImageClassification.from_pretrained('carbon225/vit-base-patch16-224-hentai')
# feature_extractor = AutoFeatureExtractor.from_pretrained('carbon225/vit-base-patch16-224-hentai')

def load_model():  
  api_key=os.getenv("GOOGLE_API_KEY")
  genai.configure(api_key=api_key)
  model = ChatGoogleGenerativeAI(model="gemini-pro",
                             temperature=0.3)
  
  return model

def load_embeddings():
  embeddings = GoogleGenerativeAIEmbeddings(model = "models/embedding-001")

  return embeddings

@app.route("/", methods=["GET"])
def default():
    return json.dumps({"Server": "Working"})

@app.route("/extractimages",methods=["GET"])
def extract_images():
    try:
        src=request.args.get("src")
        response = requests.get(src)
        soup = BeautifulSoup(response.content,'html.parser')
        img_urls=[]

        img_tags = soup.select('div img')
        for img_tag in img_tags:
            img_url = urljoin(src, img_tag['src'])
            img_urls.append(img_url)
        return json.dumps({"images":img_urls})
    except Exception as e:
        return e

# @app.route("/predict", methods=["GET"])
# def predict():
#     try:
#         src = request.args.get("src")

#         # Download image from the provided URL
#         response = requests.get(src)
#         response.raise_for_status()  

#         # Open and preprocess the image
#         image = Image.open(BytesIO(response.content))
#         image = image.resize((128, 128))

#         # Extract features using the pre-trained feature extractor
#         encoding = feature_extractor(images=image.convert("RGB"), return_tensors="pt")

#         # Make a prediction using the pre-trained model
#         with torch.no_grad():
#             outputs = model(**encoding)
#             logits = outputs.logits

#         # Get the predicted class index and label
#         predicted_class_idx = logits.argmax(-1).item()
#         predicted_class_label = model.config.id2label[predicted_class_idx]

#         # Return the predictions
#         return json.dumps({"class": predicted_class_label})

#     except requests.exceptions.RequestException as e:
#         return json.dumps({"error": f"Request error: {str(e)}"})
#     except Exception as e:
#         return json.dumps({"error": f"An unexpected error occurred: {str(e)}"})
    
@app.route('/answer',methods=['POST'])
def answer():
  api_key=os.getenv("GOOGLE_API_KEY")
  genai.configure(api_key=api_key)
  model=genai.GenerativeModel('gemini-pro')
  query=request.get_json()['query']
  final_query=f"""
  Following are negative reviews about my products, suggest what are the key issues from the customer feedback:{query}
  """
  response = model.generate_content(final_query)
  return json.dumps({"message":response.text})

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