import sys
import selectors
import asyncio
import torch
import clip
from PIL import Image
import gradio as gr

# =====================================================================
# 1. PYTHON 3.13 + GRADIO HOT RELOAD PATCH
# =====================================================================
# Prevent 'ValueError: Invalid file descriptor: -1' during dev reloads.
_orig_fileobj_to_fd = selectors._fileobj_to_fd
def _safe_fileobj_to_fd(fileobj):
    try:
        return _orig_fileobj_to_fd(fileobj)
    except ValueError as e:
        if "Invalid file descriptor: -1" in str(e):
            return -1
        raise e
selectors._fileobj_to_fd = _safe_fileobj_to_fd

def cleanup_filter(exctype, value, traceback):
    if issubclass(exctype, ValueError) and "Invalid file descriptor" in str(value):
        return
    sys.__excepthook__(exctype, value, traceback)
sys.excepthook = cleanup_filter
# =====================================================================

# 2. MODEL INITIALIZATION
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)

# 3. MAIN EVALUATION LOGIC
def hotornot(image, gender):
    if image is None:
        return "N/A", "N/A", "N/A", "N/A"
        
    # Safely convert Gradio NumPy array back to PIL Image
    image_pil = Image.fromarray(image.astype("uint8"), "RGB")
    image_tensor = preprocess(image_pil).unsqueeze(0).to(device)

    # Flattened text prompt structure for single-pass GPU batching
    all_terms = [
        f'a hot {gender}', f'a gross {gender}',
        f'a beautiful {gender}', f'an ugly {gender}',
        f'an attractive {gender}', f'a hideous {gender}'
    ]
    
    text_tokens = clip.tokenize(all_terms).to(device)

    with torch.no_grad():
        # Complete all evaluations simultaneously in a single forward pass
        logits_per_image, _ = model(image_tensor, text_tokens)
        probs = logits_per_image.softmax(dim=-1).cpu().numpy()[0]

    # Deconstruct batch pairs
    p1, n1 = probs[0], probs[1]
    p2, n2 = probs[2], probs[3]
    p3, n3 = probs[4], probs[5]

    # Math safely normalized between 0 and 100
    hotness_score = round(((p1 - n1) + 1) * 50, 2)
    beauty_score = round(((p2 - n2) + 1) * 50, 2)
    attractiveness_score = round(((p3 - n3) + 1) * 50, 2)

    # Compute a balanced composite score across all traits
    avg_positive = (p1 + p2 + p3) / 3
    avg_negative = (n1 + n2 + n3) / 3
    composite = round(((avg_positive - avg_negative) + 1) * 50, 2)
    
    return composite, hotness_score, beauty_score, attractiveness_score

# 4. GRADIO INTERFACE CONFIGURATION
iface = gr.Interface(
    fn=hotornot,
    inputs=[
        gr.Image(label="Image"),
        gr.Dropdown(
            choices=['person', 'man', 'woman'],
            value='person',
            label="Gender/Identity Type"
        )
    ],
    outputs=[
        gr.Textbox(label="Total Hot or Not™ Score"),
        gr.Textbox(label="Hotness Score"),
        gr.Textbox(label="Beauty Score"),
        gr.Textbox(label="Attractiveness Score"),
    ],
    title="Hot or Not",
    description="A simple hot or not app using OpenAI's CLIP model. The input image is passed to CLIP and evaluated against relative contrasting semantic descriptions.",
)

if __name__ == "__main__":
    # Explicitly enforce share=False to mitigate proxy collisions on local SSR setups
    iface.launch(share=False)