import gradio as gr
import pandas as pd
import numpy as np
import torch
from PIL import Image
from transformers import CLIPModel, CLIPProcessor

# ─── Model ─────────────────────────────────────────────────────────────────────
MODEL_ID = "openai/clip-vit-base-patch32"
device   = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Loading CLIP on {device}…")
model     = CLIPModel.from_pretrained(MODEL_ID).to(device).eval()
processor = CLIPProcessor.from_pretrained(MODEL_ID)
print("CLIP ready")

# ─── Dataset ───────────────────────────────────────────────────────────────────
df = pd.read_parquet("steam_games_clustered_production.parquet")
embeddings_matrix = np.stack(df["embedding"].values).astype(np.float32)
norms = np.linalg.norm(embeddings_matrix, axis=1, keepdims=True)
norms[norms == 0] = 1e-9
embeddings_matrix_normed = embeddings_matrix / norms


# ─── Inference ─────────────────────────────────────────────────────────────────
@torch.no_grad()
def embed_image(image: Image.Image) -> np.ndarray:
    if image.mode != "RGB":
        image = image.convert("RGB")
    inputs  = processor(images=image, return_tensors="pt").to(device)
    outputs = model.vision_model(**inputs)
    feats   = outputs.pooler_output
    feats   = feats / feats.norm(dim=-1, keepdim=True)
    return feats[0].cpu().numpy().astype(np.float32)


def cosine_search(query_vec: np.ndarray, top_k: int) -> pd.DataFrame:
    sims = embeddings_matrix_normed @ query_vec
    idx  = np.argsort(sims)[::-1][:top_k]
    out  = df.iloc[idx].copy()
    out["similarity"] = sims[idx]
    return out


# ─── HTML rendering of results ─────────────────────────────────────────────────
def build_row(rank: int, row) -> str:
    genres = (", ".join(row["genres"])
              if hasattr(row["genres"], "__iter__") and not isinstance(row["genres"], str)
              else row["genres"])
    sim_pct    = f"{row['similarity']*100:.0f}%"
    img_url    = row["header_image"]
    name_safe  = row["name"].replace("<", "&lt;").replace(">", "&gt;")
    search_url = f"https://store.steampowered.com/search/?term={row['name'].replace(' ', '+')}"
    return f"""
<a class="sgu-row" href="{search_url}" target="_blank" rel="noopener">
  <div class="sgu-num">{rank:02d}</div>
  <div class="sgu-thumb"><img src="{img_url}" alt="{name_safe}" loading="lazy"/></div>
  <div class="sgu-info">
    <h3 class="sgu-title">{name_safe}</h3>
    <p class="sgu-genre">{genres}</p>
  </div>
  <div class="sgu-match">
    <div class="sgu-match-pct">{sim_pct}</div>
    <div class="sgu-match-label">match</div>
  </div>
</a>"""


def search_games(query_image, top_k: int):
    if query_image is None:
        return '<div class="sgu-results-wrap"><p class="sgu-empty">Upload a game image above to begin exploring.</p></div>'

    try:
        query_vec = embed_image(query_image)
    except Exception as e:
        return f'<div class="sgu-results-wrap"><p class="sgu-empty">Error: {e}</p></div>'

    results = cosine_search(query_vec, top_k=int(top_k))
    rows = "".join(build_row(i + 1, r) for i, (_, r) in enumerate(results.iterrows()))

    return f"""
<div class="sgu-results-wrap">
  <h2 class="sgu-results-header">Your Universe</h2>
  <p class="sgu-results-sub">{int(top_k)} games sharing the visual DNA of your upload</p>
  <div class="sgu-list">{rows}</div>
</div>
"""


# ─── Styling ───────────────────────────────────────────────────────────────────
APP_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800;900&display=swap');

/* ─── flood mint everywhere ─── */
html, body, .gradio-container, gradio-app, .main, .app, .contain {
  background: #5DDBC7 !important;
  font-family: 'Inter', sans-serif !important;
  color: #0a1628;
}
.gradio-container { max-width: 100% !important; padding: 0 !important; margin: 0 !important; }
footer, .footer, #footer { display: none !important; }

/* ─── HERO ─── */
.sgu-hero {
  text-align: center;
  padding: 70px 24px 32px;
  max-width: 760px;
  margin: 0 auto;
}
.sgu-headline {
  font-weight: 900;
  font-size: clamp(4rem, 11vw, 8rem);
  line-height: 0.92;
  letter-spacing: -0.04em;
  color: #0a1628;
  margin: 0 0 28px;
}
.sgu-tagline {
  font-weight: 500;
  font-size: 1.05rem;
  color: #0a1628;
  line-height: 1.55;
  margin: 0 0 22px;
}
.sgu-explainer {
  font-weight: 400;
  font-size: 0.92rem;
  color: #0a1628;
  opacity: 0.72;
  line-height: 1.65;
  max-width: 560px;
  margin: 0 auto 0;
}

/* ─── CIRCLE: native gr.Image, sized & shaped via outer wrapper ─── */
#circle-zone {
  display: flex !important;
  justify-content: center !important;
  align-items: center !important;
  margin: 50px auto 30px !important;
  padding: 0 !important;
  width: 320px !important;
  height: 320px !important;
  flex: none !important;
  min-width: 0 !important;
}
#circle-zone > div,
#circle-zone .block,
#circle-zone .form,
#circle-zone .gradio-image,
#circle-zone .image-container {
  width: 320px !important;
  height: 320px !important;
  min-width: 0 !important;
  max-width: 320px !important;
  border-radius: 50% !important;
  overflow: hidden !important;
  background: #2D7DD2 !important;
  border: none !important;
  box-shadow: 0 20px 60px -15px rgba(45,125,210,.45);
  padding: 0 !important;
  margin: 0 auto !important;
}
#circle-zone label,
#circle-zone .label-wrap { display: none !important; }

/* The drop area inside */
#circle-zone .upload-container,
#circle-zone [data-testid="image"],
#circle-zone .wrap {
  width: 100% !important;
  height: 100% !important;
  background: #2D7DD2 !important;
  border: none !important;
  border-radius: 50% !important;
  display: flex !important;
  align-items: center !important;
  justify-content: center !important;
  cursor: pointer !important;
  color: #fff !important;
}

/* Hide Gradio's default upload icon/text, replace with custom prompt */
#circle-zone .upload-text,
#circle-zone .or,
#circle-zone svg { display: none !important; }

#circle-zone .upload-container::after,
#circle-zone [data-testid="image"]::after {
  content: "Upload a photo of your favorite game";
  position: absolute;
  text-align: center;
  font-family: 'Inter', sans-serif;
  font-weight: 500;
  font-size: 1rem;
  line-height: 1.5;
  color: #ffffff;
  width: 70%;
  pointer-events: none;
}
/* Hide that text once an image is loaded */
#circle-zone:has(img)::after,
#circle-zone [data-testid="image"]:has(img)::after { display: none !important; }

/* Make uploaded image fill the circle */
#circle-zone img {
  width: 100% !important;
  height: 100% !important;
  object-fit: cover !important;
  border-radius: 50% !important;
  display: block !important;
}

/* ─── SLIDER ─── */
#slider-zone {
  max-width: 380px !important;
  margin: 0 auto !important;
  padding: 8px 24px 0 !important;
}
#slider-zone label,
#slider-zone .label-wrap span {
  color: #0a1628 !important;
  font-weight: 500 !important;
  font-size: 0.7rem !important;
  letter-spacing: 0.18em !important;
  text-transform: uppercase !important;
}
#slider-zone input[type="range"] { accent-color: #0a1628 !important; }
#slider-zone .block { background: transparent !important; border: none !important; padding: 0 !important; }

/* ─── EXPLORE BUTTON ─── */
#explore-zone {
  display: flex !important;
  justify-content: center !important;
  padding: 28px 24px 60px !important;
  background: transparent !important;
}
#explore-zone button, #explore-btn {
  background: #0a1628 !important;
  color: #5DDBC7 !important;
  font-family: 'Inter', sans-serif !important;
  font-weight: 700 !important;
  font-size: 0.78rem !important;
  letter-spacing: 0.28em !important;
  text-transform: uppercase !important;
  padding: 18px 56px !important;
  border-radius: 999px !important;
  border: none !important;
  cursor: pointer !important;
  transition: transform .2s ease, background .2s ease !important;
  box-shadow: 0 10px 30px -10px rgba(10,22,40,.35);
}
#explore-zone button:hover, #explore-btn:hover {
  transform: translateY(-2px) !important;
  background: #163251 !important;
}

/* ─── RESULTS ─── */
.sgu-results-wrap {
  max-width: 920px;
  margin: 0 auto;
  padding: 30px 24px 80px;
  animation: sguFade .5s ease-out;
}
@keyframes sguFade { from {opacity:0;transform:translateY(16px);} to {opacity:1;transform:none;} }

.sgu-results-header {
  font-weight: 900;
  font-size: 2.2rem;
  letter-spacing: -0.02em;
  color: #0a1628;
  text-align: center;
  margin: 0 0 6px;
}
.sgu-results-sub {
  font-weight: 400;
  font-size: 0.9rem;
  color: #0a1628;
  opacity: 0.6;
  text-align: center;
  margin: 0 0 48px;
}

.sgu-list { display: flex; flex-direction: column; gap: 14px; }

.sgu-row {
  display: grid;
  grid-template-columns: 56px 180px 1fr auto;
  align-items: center;
  gap: 22px;
  background: rgba(255, 255, 255, 0.55);
  backdrop-filter: blur(8px);
  -webkit-backdrop-filter: blur(8px);
  border: 1px solid rgba(10,22,40,.08);
  border-radius: 18px;
  padding: 16px 22px;
  text-decoration: none;
  color: inherit;
  transition: transform .25s cubic-bezier(.2,.8,.2,1), box-shadow .25s ease, background .25s ease;
}
.sgu-row:hover {
  transform: translateY(-3px);
  background: rgba(255,255,255,0.85);
  box-shadow: 0 18px 50px -20px rgba(10,22,40,.25);
}

.sgu-num {
  font-weight: 900;
  font-size: 2.2rem;
  color: #0a1628;
  line-height: 1;
  letter-spacing: -0.04em;
  font-variant-numeric: tabular-nums;
  text-align: center;
}

.sgu-thumb {
  width: 180px; height: 84px;
  border-radius: 10px;
  overflow: hidden;
  background: #0a1628;
}
.sgu-thumb img {
  width: 100%; height: 100%; object-fit: cover; display: block;
  transition: transform .35s ease;
}
.sgu-row:hover .sgu-thumb img { transform: scale(1.06); }

.sgu-info { min-width: 0; }
.sgu-title {
  font-weight: 700; font-size: 1.1rem;
  color: #0a1628;
  margin: 0 0 4px;
  letter-spacing: -0.01em;
  white-space: nowrap; overflow: hidden; text-overflow: ellipsis;
}
.sgu-genre {
  font-weight: 400; font-size: 0.76rem;
  color: #0a1628; opacity: 0.55;
  margin: 0;
  white-space: nowrap; overflow: hidden; text-overflow: ellipsis;
}

.sgu-match { text-align: right; }
.sgu-match-pct {
  font-weight: 700; font-size: 1.3rem;
  color: #1976d2;
  letter-spacing: -0.02em;
  line-height: 1;
  font-variant-numeric: tabular-nums;
}
.sgu-match-label {
  font-weight: 400; font-size: 0.62rem;
  color: #0a1628; opacity: 0.5;
  text-transform: uppercase;
  letter-spacing: 0.18em;
  margin-top: 4px;
}

.sgu-empty {
  text-align: center;
  color: #0a1628; opacity: 0.55;
  padding: 60px 20px;
  font-size: 0.95rem;
}

@media (max-width: 720px) {
  .sgu-row { grid-template-columns: 40px 100px 1fr; gap: 12px; padding: 12px 14px; }
  .sgu-num { font-size: 1.6rem; }
  .sgu-thumb { width: 100px; height: 48px; }
  .sgu-title { font-size: .92rem; }
  .sgu-match { display: none; }
  .sgu-results-header { font-size: 1.7rem; }
}
"""


# ─── UI ────────────────────────────────────────────────────────────────────────
with gr.Blocks(
    theme=gr.themes.Base(),
    title="Steam Graphic Universe",
    css=APP_CSS,
    analytics_enabled=False,
) as demo:

    gr.HTML("""
    <div class="sgu-hero">
      <h1 class="sgu-headline">Let's<br/>Play!</h1>
      <p class="sgu-tagline">
        You decide the games you want to play<br/>
        We are here to help you make the best decision
      </p>
      <p class="sgu-explainer">
        Upload an image of any game and we'll surface the closest visual matches
        from a universe of over 9,800 Steam titles. Behind the scenes a CLIP vision
        model converts your image into a high-dimensional aesthetic fingerprint,
        then ranks every game by visual similarity — no genres, no keywords,
        just pure art-style resonance.
      </p>
    </div>
    """)

    # Image uploader — let Gradio handle the upload natively, we just shape the wrapper
    query_image = gr.Image(
        type="pil",
        show_label=False,
        show_download_button=False,
        sources=["upload", "clipboard"],
        elem_id="circle-zone",
        height=320,
        width=320,
    )

    top_k_slider = gr.Slider(
        minimum=3, maximum=12, step=1, value=6,
        label="Number of matches",
        elem_id="slider-zone",
    )

    with gr.Row(elem_id="explore-zone"):
        search_btn = gr.Button("Explore", elem_id="explore-btn")

    output_html = gr.HTML(
        value='<div class="sgu-results-wrap"><p class="sgu-empty">Upload a game image above to begin exploring.</p></div>'
    )

    search_btn.click(
        fn=search_games,
        inputs=[query_image, top_k_slider],
        outputs=output_html,
    )

if __name__ == "__main__":
    demo.launch()