import gradio as gr
import requests
from bs4 import BeautifulSoup
from PIL import Image
from transformers import ViltProcessor, ViltForQuestionAnswering
import torch
import torch.nn.functional as F

# ViLT Model betöltése
model_path = "dandelin/vilt-b32-finetuned-vqa"  # Hugging Face modell
device = "cuda" if torch.cuda.is_available() else "cpu"

model = ViltForQuestionAnswering.from_pretrained(model_path).to(device)
processor = ViltProcessor.from_pretrained(model_path)

# Keresés a Roblox katalógusban
def search_roblox(character_name):
    url = f"https://www.roblox.com/catalog?Category=3&salesTypeFilter=1"
    headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.36'}
    response = requests.get(url, headers=headers)
    
    if response.status_code != 200:
        return None, "Failed to fetch data", 0
    
    soup = BeautifulSoup(response.text, "html.parser")
    
    # Az elemzéshez szükséges adatokat a megfelelő HTML elemekből szűrjük
    items = soup.find_all('div', {'class': 'catalog-item-container'})
    
    best_item = None
    best_score = 0

    for item in items[:5]:  # Csak az első 5 találatot nézzük
        try:
            img_tag = item.find('img')
            img_url = img_tag['src'] if img_tag else None
            name = item.find('div', {'class': 'item-card-name'}).text.strip()
            link_tag = item.find('a', {'class': 'item-card-link'})
            link = link_tag['href'] if link_tag else None
            
            if not img_url or not link:
                continue

            image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")

            # AI kiértékelés
            question = f"Is this item good for the character: {character_name}? Give a score from 0 to 1000."
            inputs = processor(images=image, text=question, return_tensors="pt").to(device)
            outputs = model(**inputs)

            softmax_scores = F.softmax(outputs.logits, dim=-1)
            score = softmax_scores.max().item() * 1000  # 0-1000 közötti skála

            if score > best_score:
                best_score = score
                best_item = {"name": name, "img_url": img_url, "link": link, "score": best_score}

        except Exception as e:
            print(f"Error processing item: {e}")

    if best_item:
        return best_item["img_url"], best_item["link"], best_item["score"]
    else:
        return None, "No items found", 0

# Gradio UI
iface = gr.Interface(
    fn=search_roblox,
    inputs=gr.Textbox(label="Enter Character Name"),
    outputs=[gr.Image(label="Best Match"), gr.Textbox(label="Item Link"), gr.Number(label="AI Score")],
    title="Roblox Character Item Finder",
    description="Enter a character name, and the AI will find the best Roblox catalog item for it."
)

iface.launch()