import os, base64, requests, pickle
import pandas as pd
import numpy as np
import torch
import torchvision.transforms as T
from torchvision.models import resnet50, ResNet50_Weights
from sklearn.metrics.pairwise import cosine_similarity
from PIL import Image
import gradio as gr
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
API_KEY = os.getenv("ROBOFLOW_API_KEY")
MODEL_ENDPOINT = "product-fashion-matching-02/2"

# Load model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
resnet = resnet50(weights=ResNet50_Weights.DEFAULT)
resnet = torch.nn.Sequential(*list(resnet.children())[:-1])
resnet.eval().to(device)

transform = T.Compose([
    T.Resize((224, 224)),
    T.ToTensor(),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

# Load precomputed features
with open("models/similarity_model.pkl", "rb") as f:
    data = pickle.load(f)
features_np = data["features"]
product_ids = data["product_ids"]
df = pd.read_csv("data/ready_dataset.csv")

def crop_with_bbox(pil_img):
    temp_path = "static/temp.jpg"
    pil_img.save(temp_path)

    with open(temp_path, "rb") as f:
        img_b64 = base64.b64encode(f.read()).decode()

    url = f"https://detect.roboflow.com/{MODEL_ENDPOINT}"
    params = {"api_key": API_KEY}
    headers = {"Content-Type": "application/x-www-form-urlencoded"}

    r = requests.post(url, params=params, data=img_b64, headers=headers)
    r.raise_for_status()
    preds = r.json().get("predictions", [])
    if not preds:
        return None

    p = preds[0]
    x, y, w, h = p['x'], p['y'], p['width'], p['height']
    x1, y1 = int(x - w/2), int(y - h/2)
    x2, y2 = int(x + w/2), int(y + h/2)

    return pil_img.crop((x1, y1, x2, y2))

def predict_similar_products(image):
    if image is None:
        return "❌ Please upload an image.", [], []

    cropped = crop_with_bbox(image)
    if cropped is None:
        return "❌ No bounding box detected.", [], []

    tensor = transform(cropped).unsqueeze(0).to(device)
    with torch.no_grad():
        qf = resnet(tensor).squeeze().cpu().numpy()

    sims = cosine_similarity([qf], features_np)[0]
    top_idxs = sims.argsort()[-10:][::-1]
    top_pids = [product_ids[i] for i in top_idxs]

    res = df[df.product_id.isin(top_pids)][["product_name", "feature_image_s3"]].drop_duplicates()
    names = res["product_name"].tolist()
    urls = res["feature_image_s3"].tolist()

    return "✅ Top 10 Similar Products Found:", names, urls

# Custom CSS (injected via Markdown)
custom_css = """
<style>
body {
  background-color: #f5f5f5;
  font-family: 'Segoe UI', sans-serif;
  text-align: center;
  padding-top: 20px;
}

h1, h2 {
  color: #2c3e50;
  font-weight: bold;
}

.gradio-container {
  max-width: 1200px;
  margin: 0 auto;
}

.gr-box, .gr-panel {
  background-color: #ffffff !important;
  border-radius: 16px;
  padding: 24px;
  box-shadow: 0 10px 20px rgba(0,0,0,0.07);
  margin-bottom: 20px;
}

button {
  background-color: #3498db !important;
  color: white !important;
  padding: 12px 24px !important;
  font-weight: 600 !important;
  border-radius: 8px !important;
  border: none !important;
  cursor: pointer !important;
}

button:hover {
  background-color: #2c80b4 !important;
}

img {
  border-radius: 12px;
  object-fit: cover;
  max-height: 300px;
  width: auto;
}

#gallery {
  display: grid;
  grid-template-columns: repeat(auto-fill, minmax(180px, 1fr));
  gap: 20px;
  justify-items: center;
  padding: 20px;
}

#gallery > div {
  background-color: #ffffff;
  padding: 10px;
  border-radius: 12px;
  box-shadow: 0 5px 10px rgba(0,0,0,0.1);
  transition: transform 0.2s ease;
}

#gallery > div:hover {
  transform: scale(1.05);
}
</style>
"""

with gr.Blocks() as demo:
    gr.Markdown(custom_css)
    gr.Markdown("## 🧥 Fashion Product Similarity")
    gr.Markdown("Upload a fashion product image. We'll detect the Top 10 Closest matches of your image according to the dataset.")
    
    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Fashion Product Image")

    with gr.Row():
        submit_btn = gr.Button("🔍 Find Similar Products")

    status_output = gr.Textbox(label="Status")
    name_output = gr.Textbox(label="Top Product Names")
    gallery_output = gr.Gallery(label="Top 10 Similar Products", columns=5, elem_id="gallery")

    submit_btn.click(
        fn=predict_similar_products,
        inputs=image_input,
        outputs=[status_output, name_output, gallery_output]
    )

if __name__ == "__main__":
    demo.launch(share=True, ssr_mode=False)