app.py

import os
import io
import json
import time
import shutil
import random
from datetime import datetime

from flask import Flask, request, jsonify, render_template
from huggingface_hub import HfApi, hf_hub_download
from PIL import Image

import numpy as np
import faiss
import torch
from transformers import CLIPProcessor, CLIPModel

# -------------------
# Config
# -------------------
HF_TOKEN = os.getenv("HF_TOKEN")  # put your real token in HF Space Secrets
DATASET_REPO = "Sahil5112/Fan_hub"  # your dataset repo

IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".gif", ".webp")
VIDEO_EXTS = (".mp4", ".webm", ".mov", ".avi")

FAISS_FILENAME = "index.faiss"
MAP_FILENAME = "index_map.json"
META_FILENAME = "metadata.json"

STATE_DIR = "state"
UPLOADS_DIR = "uploads"

os.makedirs(STATE_DIR, exist_ok=True)
os.makedirs(UPLOADS_DIR, exist_ok=True)

app = Flask(__name__, template_folder="templates")
api = HfApi()

# -------------------
# CLIP Model
# -------------------
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MODEL_ID = "openai/clip-vit-base-patch32"
clip_model = CLIPModel.from_pretrained(MODEL_ID).to(DEVICE)
clip_processor = CLIPProcessor.from_pretrained(MODEL_ID)
EMB_DIM = clip_model.visual_projection.out_features

# -------------------
# Helpers
# -------------------
def l2_normalize(x: np.ndarray) -> np.ndarray:
    x = x.astype("float32")
    norms = np.linalg.norm(x, axis=1, keepdims=True) + 1e-10
    return x / norms

def image_to_embedding(pil_img: Image.Image) -> np.ndarray:
    if pil_img.mode != "RGB":
        pil_img = pil_img.convert("RGB")
    pil_img.thumbnail((512, 512), Image.Resampling.LANCZOS)
    inputs = clip_processor(images=pil_img, return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        feats = clip_model.get_image_features(**inputs)
    feats = feats.cpu().numpy().astype("float32")
    return l2_normalize(feats)

def text_to_embedding(text: str) -> np.ndarray:
    inputs = clip_processor(text=[text], return_tensors="pt", padding=True).to(DEVICE)
    with torch.no_grad():
        feats = clip_model.get_text_features(**inputs)
    feats = feats.cpu().numpy().astype("float32")
    return l2_normalize(feats)

def dataset_url(path_in_repo: str) -> str:
    return f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{path_in_repo}"

def load_json_local(path: str, default):
    if os.path.exists(path):
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    return default

def save_json_local(path: str, obj):
    with open(path, "w", encoding="utf-8") as f:
        json.dump(obj, f, ensure_ascii=False, indent=2)

# -------------------
# Persistent Index
# -------------------
index: faiss.Index = None
index_map: list[dict] = []
metadata: dict = {}

def _create_empty_index():
    return faiss.IndexFlatIP(EMB_DIM)

def pull_or_init_state():
    global index, index_map, metadata
    for repo_path in [FAISS_FILENAME, MAP_FILENAME, META_FILENAME]:
        try:
            hf_hub_download(
                repo_id=DATASET_REPO,
                filename=repo_path,
                repo_type="dataset",
                token=HF_TOKEN,
                local_dir=STATE_DIR,
                local_dir_use_symlinks=False
            )
        except Exception:
            pass

    faiss_path = os.path.join(STATE_DIR, FAISS_FILENAME)
    map_path = os.path.join(STATE_DIR, MAP_FILENAME)
    meta_path = os.path.join(STATE_DIR, META_FILENAME)

    index = faiss.read_index(faiss_path) if os.path.exists(faiss_path) else _create_empty_index()
    index_map = load_json_local(map_path, [])
    metadata = load_json_local(meta_path, {})

    for fn, md in metadata.items():
        if "uploader" not in md:
            md["uploader"] = "Anonymous"
    save_json_local(meta_path, metadata)

def persist_state():
    faiss_path = os.path.join(STATE_DIR, FAISS_FILENAME)
    map_path = os.path.join(STATE_DIR, MAP_FILENAME)
    meta_path = os.path.join(STATE_DIR, META_FILENAME)

    faiss.write_index(index, faiss_path)
    save_json_local(map_path, index_map)
    save_json_local(meta_path, metadata)

    for local_file, repo_file in [(faiss_path, FAISS_FILENAME), (map_path, MAP_FILENAME), (meta_path, META_FILENAME)]:
        api.upload_file(
            path_or_fileobj=local_file,
            path_in_repo=repo_file,
            repo_id=DATASET_REPO,
            repo_type="dataset",
            token=HF_TOKEN
        )

def add_item_to_index(filename: str, pil_img: Image.Image):
    vec = image_to_embedding(pil_img)
    index.add(vec)
    md = metadata.get(filename, {})
    index_map.append({
        "filename": filename,
        "title": md.get("title", "Untitled"),
        "description": md.get("description", ""),
        "category": md.get("category", "Uncategorized"),
        "filetype": "image"
    })

def top_k_from_index(query_vec: np.ndarray, k: int = 24, exclude_filename: str | None = None):
    if index.ntotal == 0:
        return []

    D, I = index.search(query_vec.astype("float32"), k + 1)
    results = []
    for idx in I[0]:
        if idx < 0 or idx >= len(index_map):
            continue
        item = index_map[idx]
        if exclude_filename and item["filename"] == exclude_filename:
            continue
        results.append({
            "url": metadata.get(item["filename"], {}).get("url", dataset_url(item["filename"])),
            "filename": item["filename"],
            "title": item.get("title", "Untitled"),
            "description": item.get("description", ""),
            "category": item.get("category", "Uncategorized"),
            "uploader": metadata.get(item["filename"], {}).get("uploader", "Anonymous"),
            "filetype": item.get("filetype", "image")
        })
        if len(results) >= k:
            break
    return results

# -------------------
# Routes
# -------------------
@app.route("/")
def home():
    return render_template("index.html")

@app.route("/upload", methods=["POST"])
def upload():
    if "file" not in request.files:
        return jsonify({"error": "No file provided"}), 400
    f = request.files["file"]

    ext = os.path.splitext(f.filename.lower())[1]
    if not (ext in IMAGE_EXTS or ext in VIDEO_EXTS):
        return jsonify({"error": "Unsupported file type"}), 400

    title = request.form.get("title", "Untitled").strip()[:120]
    description = request.form.get("description", "").strip()[:500]
    category = request.form.get("category", "Uncategorized").strip()[:60]
    uploader = request.form.get("uploader", "Anonymous").strip()[:80]

    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    safe_name = f.filename.replace("/", "_").replace("\\", "_")
    filename = f"{ts}_{safe_name}"
    local_path = os.path.join(UPLOADS_DIR, filename)

    if ext in IMAGE_EXTS and ext != ".gif":
        try:
            img = Image.open(f).convert("RGB")
            img.save(local_path, format="JPEG", optimize=True, quality=85)
        except Exception as e:
            print("Compression failed:", e)
            f.seek(0)
            f.save(local_path)
    else:
        f.save(local_path)

    api.upload_file(
        path_or_fileobj=local_path,
        path_in_repo=filename,
        repo_id=DATASET_REPO,
        repo_type="dataset",
        token=HF_TOKEN
    )

    filetype = "video" if ext in VIDEO_EXTS else "image"
    metadata[filename] = {
        "title": title or "Untitled",
        "description": description,
        "category": category or "Uncategorized",
        "uploader": uploader or "Anonymous",
        "filetype": filetype,
        "url": dataset_url(filename)
    }

    if filetype == "image" and ext != ".gif":
        img = Image.open(local_path).convert("RGB")
        add_item_to_index(filename=filename, pil_img=img)

    persist_state()
    try: os.remove(local_path)
    except: pass

    return jsonify({
        "success": True,
        "filename": filename,
        "url": metadata[filename]["url"],
        "filetype": filetype,
        "uploader": metadata[filename]["uploader"]
    })

# -------------------
# Gallery with Pagination (Lazy Load)
# -------------------
@app.route("/gallery")
def gallery():
    try:
        offset = int(request.args.get("offset", 0))
        limit = int(request.args.get("limit", 10))
    except:
        offset = 0
        limit = 10

    # --- Get filters from query params ---
    text_filter = request.args.get("text", "").lower()
    uploader_filter = request.args.get("uploader", "").lower()
    category_filter = request.args.get("category", "").lower()
    filetype_filter = request.args.get("filetype", "").lower()
    tags_filter = [t.strip().lower() for t in request.args.get("tags", "").split(",") if t.strip()]

    items = []

    # Images from index_map
    for entry in index_map:
        fn = entry["filename"]
        md = metadata.get(fn, {})
        item = {
            "url": md.get("url", dataset_url(fn)),
            "filename": fn,
            "title": md.get("title", entry.get("title", "Untitled")),
            "description": md.get("description", entry.get("description", "")),
            "category": md.get("category", entry.get("category", "Uncategorized")),
            "uploader": md.get("uploader", "Anonymous"),
            "filetype": md.get("filetype", entry.get("filetype", "image")),
            "tags": md.get("tags", [])
        }

        # --- Apply filters ---
        if text_filter and text_filter not in item["title"].lower() and text_filter not in item["description"].lower():
            continue
        if uploader_filter and uploader_filter != item["uploader"].lower():
            continue
        if category_filter and category_filter != item["category"].lower():
            continue
        if filetype_filter and filetype_filter != item["filetype"].lower():
            continue
        if tags_filter and not any(t.lower() in [tag.lower() for tag in item.get("tags", [])] for t in tags_filter):
            continue

        items.append(item)

    # Add videos not in index_map
    for fn, md in metadata.items():
        if md.get("filetype") == "video" and not any(it["filename"] == fn for it in items):
            item = {
                "url": md.get("url", dataset_url(fn)),
                "filename": fn,
                "title": md.get("title", "Untitled"),
                "description": md.get("description", ""),
                "category": md.get("category", "Uncategorized"),
                "uploader": md.get("uploader", "Anonymous"),
                "filetype": "video",
                "tags": md.get("tags", [])
            }

            # --- Apply filters ---
            if text_filter and text_filter not in item["title"].lower() and text_filter not in item["description"].lower():
                continue
            if uploader_filter and uploader_filter != item["uploader"].lower():
                continue
            if category_filter and category_filter != item["category"].lower():
                continue
            if filetype_filter and filetype_filter != item["filetype"].lower():
                continue
            if tags_filter and not any(t.lower() in [tag.lower() for tag in item.get("tags", [])] for t in tags_filter):
                continue

            items.append(item)

    items.sort(key=lambda x: x["filename"], reverse=True)
    paged_items = items[offset: offset + limit]
    return jsonify(paged_items)

# -------------------
# Related & Search
# -------------------
@app.route("/related", methods=["POST"])
def related():
    data = request.get_json(force=True)
    filename = data.get("filename", "")
    if not filename:
        return jsonify([])

    md = metadata.get(filename, {})
    filetype = md.get("filetype", "image")

    if filetype == "image":
        try:
            local = hf_hub_download(repo_id=DATASET_REPO, filename=filename, repo_type="dataset", token=HF_TOKEN)
            img = Image.open(local).convert("RGB")
            qvec = image_to_embedding(img)
            return jsonify(top_k_from_index(qvec, k=12, exclude_filename=filename))
        except Exception:
            return jsonify([])
    else:
        same_cat = [fn for fn, m in metadata.items() if fn != filename and m.get("category") == md.get("category")]
        random.shuffle(same_cat)
        res = []
        for fn in same_cat[:12]:
            m = metadata[fn]
            res.append({
                "url": m.get("url", dataset_url(fn)),
                "filename": fn,
                "title": m.get("title", "Untitled"),
                "description": m.get("description", ""),
                "category": m.get("category", "Uncategorized"),
                "uploader": m.get("uploader", "Anonymous"),
                "filetype": m.get("filetype", "image")
            })
        return jsonify(res)

@app.route("/search/text")
def search_text():
    q = request.args.get("q", "").strip()
    if not q:
        return jsonify([])
    qvec = text_to_embedding(q)
    return jsonify(top_k_from_index(qvec, k=30))

@app.route("/search/image", methods=["POST"])
def search_image():
    if "file" not in request.files:
        return jsonify({"error": "No file"}), 400
    f = request.files["file"]
    if not f.filename.lower().endswith(IMAGE_EXTS):
        return jsonify({"error": "Unsupported file type"}), 400

    img = Image.open(io.BytesIO(f.read())).convert("RGB")
    qvec = image_to_embedding(img)
    return jsonify(top_k_from_index(qvec, k=30))

# -------------------
# Init
# -------------------
pull_or_init_state()

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=False)