"""
app.py  —  Fashion Recommendation API  (CLIP + FAISS + Re-ranking)
"""

import os
import io
import pickle
import tempfile
import requests
import numpy as np
import pandas as pd
from PIL import Image
from contextlib import asynccontextmanager

from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

import torch
import open_clip
import faiss

from dotenv import load_dotenv

load_dotenv()

# ─────────────────────────────────────────────
# ENV CONFIG
# ─────────────────────────────────────────────
HF_OWNER         = os.getenv("HF_DATASET_OWNER")
HF_REPO          = os.getenv("HF_DATASET_REPO")
CLOUDINARY_CLOUD = os.getenv("CLOUDINARY_CLOUD_NAME")

if not all([HF_OWNER, HF_REPO, CLOUDINARY_CLOUD]):
    raise ValueError("❌ Missing env vars: HF_DATASET_OWNER, HF_DATASET_REPO, CLOUDINARY_CLOUD_NAME")

BASE_HF        = f"https://huggingface.co/datasets/{HF_OWNER}/{HF_REPO}/resolve/main"
EMBEDDINGS_URL = f"{BASE_HF}/clip_embeddings.pkl"
FILENAMES_URL  = f"{BASE_HF}/clip_filenames.pkl"
FAISS_URL      = f"{BASE_HF}/faiss_index.bin"
CSV_URL        = f"{BASE_HF}/styles.csv"
CLOUDINARY_URL = f"https://res.cloudinary.com/{CLOUDINARY_CLOUD}/image/upload/"

CLIP_MODEL_NAME = "ViT-B-32"
CLIP_PRETRAINED = "openai"
FAISS_NPROBE    = 10
ANN_CANDIDATES  = 50
FINAL_TOP_K     = 5

ml_models = {}
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


# ─────────────────────────────────────────────
# HELPERS: Download
# ─────────────────────────────────────────────
def download_bytes(url: str) -> bytes:
    print(f"  ↓ {url}")
    r = requests.get(url, stream=True)
    r.raise_for_status()
    return r.content

def load_pickle(url: str):
    return pickle.loads(download_bytes(url))

def load_csv(url: str) -> pd.DataFrame:
    content = download_bytes(url)
    return pd.read_csv(io.BytesIO(content), on_bad_lines="skip")

def load_faiss_index(url: str) -> faiss.Index:
    data = download_bytes(url)
    with tempfile.NamedTemporaryFile(delete=False, suffix=".bin") as f:
        f.write(data)
        tmp_path = f.name
    index = faiss.read_index(tmp_path)
    os.unlink(tmp_path)
    return index


# ─────────────────────────────────────────────
# STARTUP / SHUTDOWN
# ─────────────────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    print(f"🚀 Starting up (device={DEVICE})…")
    try:
        print("Loading CLIP model…")
        clip_model, _, preprocess = open_clip.create_model_and_transforms(
            CLIP_MODEL_NAME, pretrained=CLIP_PRETRAINED, device=DEVICE
        )
        clip_model.eval()

        print("Downloading embeddings, filenames, FAISS index, styles…")
        embeddings  = np.array(load_pickle(EMBEDDINGS_URL), dtype="float32")
        filenames   = load_pickle(FILENAMES_URL)
        styles_df   = load_csv(CSV_URL)
        faiss_index = load_faiss_index(FAISS_URL)

        if hasattr(faiss_index, "nprobe"):
            faiss_index.nprobe = FAISS_NPROBE

        ml_models.update({
            "clip_model":  clip_model,
            "preprocess":  preprocess,
            "embeddings":  embeddings,
            "filenames":   filenames,
            "styles_df":   styles_df,
            "faiss_index": faiss_index,
        })
        print(f"✅ Ready! Index has {faiss_index.ntotal} vectors.")

    except Exception as e:
        print(f"🛑 Startup error: {e}")
        raise

    yield
    ml_models.clear()
    print("🛑 Shutdown.")


# ─────────────────────────────────────────────
# APP
# ─────────────────────────────────────────────
app = FastAPI(title="Klaro — Fashion Recommendation API (CLIP + FAISS)", lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True,
                   allow_methods=["*"], allow_headers=["*"])


# ─────────────────────────────────────────────
# CORE: Feature extraction
# ─────────────────────────────────────────────
@torch.no_grad()
def encode_image(img: Image.Image) -> np.ndarray:
    preprocess = ml_models["preprocess"]
    clip_model = ml_models["clip_model"]
    tensor   = preprocess(img).unsqueeze(0).to(DEVICE)
    features = clip_model.encode_image(tensor)
    features = features / features.norm(dim=-1, keepdim=True)
    return features.cpu().numpy().astype("float32")   # (1, 512)


# ─────────────────────────────────────────────
# CORE: Two-stage retrieval
# ─────────────────────────────────────────────
def faiss_retrieve(query_vec: np.ndarray, k: int):
    index = ml_models["faiss_index"]
    scores, indices = index.search(query_vec, k)
    return scores[0], indices[0]

def rerank(query_vec: np.ndarray, candidate_indices: np.ndarray, top_k: int):
    embeddings     = ml_models["embeddings"]
    candidate_embs = embeddings[candidate_indices]
    scores = (candidate_embs @ query_vec.T).flatten()
    order  = np.argsort(scores)[::-1][:top_k]
    return scores[order], candidate_indices[order]


# ─────────────────────────────────────────────
# HELPERS: Metadata  ← FIXED (ml_models + new CSV columns)
# ─────────────────────────────────────────────
def get_metadata(filename: str) -> dict:
    df = ml_models["styles_df"]          # ← was state["styles_df"], now fixed
    try:
        image_id = int(os.path.splitext(filename)[0])
        row = df[df["Product ID"] == image_id]
        if not row.empty:
            return {
                "gender":          str(row["Age Group"].values[0]),
                "articleType":     str(row["Sub Category"].values[0]),
                "baseColour":      str(row["Base Colour"].values[0]),
                "usage":           str(row["Season"].values[0]),
                "name":            str(row["Product Name"].values[0]),
                "brand":           str(row["Brand"].values[0]),
                "price":           float(row["Price"].values[0]),
                "discountedPrice": float(row["Discounted Price"].values[0]),
            }
    except Exception as e:
        print(f"Metadata error for {filename}: {e}")
    return {}


# ─────────────────────────────────────────────
# ENDPOINT
# ─────────────────────────────────────────────
@app.post("/recommend")
async def recommend_fashion(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        img = Image.open(io.BytesIO(contents)).convert("RGB")

        # Stage 1: CLIP encode + FAISS ANN
        query_vec = encode_image(img)                         # (1, 512)
        _, ann_indices = faiss_retrieve(query_vec, k=ANN_CANDIDATES)
        ann_indices = ann_indices[ann_indices >= 0]           # remove -1s

        # Stage 2: Exact cosine re-rank → top-5
        scores, top_indices = rerank(query_vec, ann_indices, top_k=FINAL_TOP_K)

        filenames = ml_models["filenames"]

        results = []
        for idx, score in zip(top_indices, scores):
            filename  = os.path.basename(filenames[idx])
            metadata  = get_metadata(filename)
            results.append({
                "filename":  filename,
                "image_url": CLOUDINARY_URL + filename,
                "metadata":  metadata,
                "score":     float(score),
            })

        return {"recommendations": results}

    except Exception as e:
        import traceback
        traceback.print_exc()
        return JSONResponse(status_code=500, content={"error": str(e)})


@app.get("/health")
def health():
    idx = ml_models.get("faiss_index")
    return {
        "status":     "ok",
        "index_size": idx.ntotal if idx else 0,
        "device":     DEVICE,
    }


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