main.py

from contextlib import asynccontextmanager
from collections import OrderedDict
import asyncio
import os
import shutil
import uuid
import re
import inflect
from urllib.parse import urlparse
from typing import List

from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware

import cloudinary
import cloudinary.uploader
import cloudinary.api
from pinecone import Pinecone, ServerlessSpec

# ── Deferred imports so startup prints appear in order ────────────
ai      = None   # set in lifespan
p       = inflect.engine()

# ── Semaphore: max concurrent AI inference jobs ────────────────────
MAX_CONCURRENT_INFERENCES = int(os.getenv("MAX_CONCURRENT_INFERENCES", "6"))
_inference_sem: asyncio.Semaphore

# ── Simple LRU connection pools ───────────────────────────────────
_pinecone_pool: OrderedDict  = OrderedDict()
_cloudinary_pool: dict = {}
_POOL_MAX = 64


def _get_pinecone(api_key: str) -> Pinecone:
    """Return a cached Pinecone client, creating one if needed."""
    if api_key not in _pinecone_pool:
        if len(_pinecone_pool) >= _POOL_MAX:
            _pinecone_pool.popitem(last=False)   # evict oldest
        _pinecone_pool[api_key] = Pinecone(api_key=api_key)
    _pinecone_pool.move_to_end(api_key)          # refresh LRU order
    return _pinecone_pool[api_key]


def _configure_cloudinary(creds: dict) -> None:
    """Configure cloudinary module only when needed, with simple caching."""
    key = creds["cloud_name"]
    if key not in _cloudinary_pool:
        cloudinary.config(
            cloud_name=creds["cloud_name"],
            api_key=creds["api_key"],
            api_secret=creds["api_secret"],
        )
        _cloudinary_pool[key] = True


# ── Lifespan: load models once at startup ─────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    global ai, _inference_sem
    from src.models import AIModelManager

    print("⏳ Loading AI models …")
    loop = asyncio.get_event_loop()
    ai = await loop.run_in_executor(None, AIModelManager)
    _inference_sem = asyncio.Semaphore(MAX_CONCURRENT_INFERENCES)
    print(f"✅ Ready! Max concurrent inference slots: {MAX_CONCURRENT_INFERENCES}")
    yield
    print("👋 Shutting down")


app = FastAPI(lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],        # tighten to your Vercel domain in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

os.makedirs("temp_uploads", exist_ok=True)


# ── Helpers ────────────────────────────────────────────────────────
def standardize_category_name(name: str) -> str:
    clean = re.sub(r'\s+', '_', name.strip().lower())
    clean = re.sub(r'[^\w]', '', clean)
    return p.singular_noun(clean) or clean


def sanitize_filename(filename: str) -> str:
    clean = re.sub(r'\s+', '_', filename)
    return re.sub(r'[^\w.\-]', '', clean)


def get_cloudinary_creds(env_url: str) -> dict:
    parsed = urlparse(env_url)
    return {
        "api_key":    parsed.username,
        "api_secret": parsed.password,
        "cloud_name": parsed.hostname,
    }


# ══════════════════════════════════════════════════════════════════
# 1. VERIFY KEYS & AUTO-BUILD INDEXES
# ══════════════════════════════════════════════════════════════════
@app.post("/api/verify-keys")
async def verify_keys(
    pinecone_key:    str = Form(""),
    cloudinary_url:  str = Form(""),
):
    if cloudinary_url:
        try:
            creds = get_cloudinary_creds(cloudinary_url)
            _configure_cloudinary(creds)
            await asyncio.to_thread(cloudinary.api.ping)
        except Exception:
            raise HTTPException(400, "Invalid Cloudinary Environment URL.")

    if pinecone_key:
        try:
            pc = _get_pinecone(pinecone_key)
            existing = {idx.name for idx in await asyncio.to_thread(pc.list_indexes)}

            tasks = []
            if "lens-objects" not in existing:
                tasks.append(asyncio.to_thread(
                    pc.create_index,
                    name="lens-objects", dimension=1536, metric="cosine",
                    spec=ServerlessSpec(cloud="aws", region="us-east-1"),
                ))
            if "lens-faces" not in existing:
                tasks.append(asyncio.to_thread(
                    pc.create_index,
                    name="lens-faces", dimension=512, metric="cosine",
                    spec=ServerlessSpec(cloud="aws", region="us-east-1"),
                ))

            if tasks:
                await asyncio.gather(*tasks)

        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(400, f"Pinecone Error: {e}")

    return {"message": "Keys verified and indexes ready!"}


# ══════════════════════════════════════════════════════════════════
# 2. UPLOAD (Cloudinary + Pinecone Only)
# ══════════════════════════════════════════════════════════════════
@app.post("/api/upload")
async def upload_new_images(
    files:               List[UploadFile] = File(...),
    folder_name:         str  = Form(...),
    detect_faces:        bool = Form(True),
    user_pinecone_key:   str  = Form(""),
    user_cloudinary_url: str  = Form(""),
):
    if not user_pinecone_key or not user_cloudinary_url:
        raise HTTPException(status_code=400, detail="Cloudinary URL and Pinecone API Key are required to upload.")

    folder    = standardize_category_name(folder_name)
    uploaded_urls = []

    cld_creds = get_cloudinary_creds(user_cloudinary_url)
    _configure_cloudinary(cld_creds)
    pc        = _get_pinecone(user_pinecone_key)
    idx_obj   = pc.Index("lens-objects")
    idx_face  = pc.Index("lens-faces")

    for file in files:
        safe_name = sanitize_filename(file.filename)
        tmp_path  = f"temp_uploads/{uuid.uuid4().hex}_{safe_name}"

        try:
            with open(tmp_path, "wb") as buf:
                shutil.copyfileobj(file.file, buf)

            # Upload image to CDN
            result    = await asyncio.to_thread(cloudinary.uploader.upload, tmp_path, folder=folder)
            image_url = result["secure_url"]
            uploaded_urls.append(image_url)

            # AI inference
            async with _inference_sem:
                vectors = await ai.process_image_async(tmp_path, is_query=False, detect_faces=detect_faces)

            # Save vectors
            face_upserts   = []
            object_upserts = []

            for v in vectors:
                vec_list = v["vector"].tolist() if hasattr(v["vector"], "tolist") else v["vector"]
                record   = {
                    "id":     str(uuid.uuid4()),
                    "values": vec_list,
                    "metadata": {"url": image_url, "folder": folder},
                }
                (face_upserts if v["type"] == "face" else object_upserts).append(record)

            # Fire both upserts concurrently
            upsert_tasks = []
            if face_upserts:
                upsert_tasks.append(asyncio.to_thread(idx_face.upsert, vectors=face_upserts))
            if object_upserts:
                upsert_tasks.append(asyncio.to_thread(idx_obj.upsert, vectors=object_upserts))
            if upsert_tasks:
                await asyncio.gather(*upsert_tasks)

        except Exception as e:
            print(f"❌ Upload error for {file.filename}: {e}")
            # Continue with the next file instead of aborting the whole batch
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)

    return {"message": "Done!", "urls": uploaded_urls}


# ══════════════════════════════════════════════════════════════════
# 3. SEARCH (Pinecone Only)
# ══════════════════════════════════════════════════════════════════
@app.post("/api/search")
async def search_database(
    file:                UploadFile = File(...),
    detect_faces:        bool = Form(True),
    user_pinecone_key:   str  = Form(""),
    user_cloudinary_url: str  = Form(""), # Kept to match frontend form payload
):
    if not user_pinecone_key:
        raise HTTPException(status_code=400, detail="Pinecone API Key is required to search.")

    safe_name = sanitize_filename(file.filename)
    tmp_path  = f"temp_uploads/query_{uuid.uuid4().hex}_{safe_name}"

    try:
        with open(tmp_path, "wb") as buf:
            shutil.copyfileobj(file.file, buf)

        # AI inference
        async with _inference_sem:
            vectors = await ai.process_image_async(tmp_path, is_query=True, detect_faces=detect_faces)

        pc       = _get_pinecone(user_pinecone_key)
        idx_obj  = pc.Index("lens-objects")
        idx_face = pc.Index("lens-faces")

        # Fire ALL vector queries in parallel
        async def _query_one(vec_dict: dict) -> list[dict]:
            vec_list   = (vec_dict["vector"].tolist() if hasattr(vec_dict["vector"], "tolist") else vec_dict["vector"])
            target_idx = idx_face if vec_dict["type"] == "face" else idx_obj
            res = await asyncio.to_thread(
                target_idx.query,
                vector=vec_list, top_k=10, include_metadata=True,
            )
            out = []
            for match in res.get("matches", []):
                caption = ("👤 Verified Identity" if vec_dict["type"] == "face" else match["metadata"].get("folder", "🎯 Object Match"))
                out.append({
                    "url":     match["metadata"].get("url", ""),
                    "score":   match["score"],
                    "caption": caption,
                })
            return out

        nested = await asyncio.gather(*[_query_one(v) for v in vectors])
        all_results = [r for sub in nested for r in sub]

        # Deduplicate, keep best score per URL
        seen: dict[str, dict] = {}
        for r in all_results:
            url = r["url"]
            if url not in seen or r["score"] > seen[url]["score"]:
                seen[url] = r

        final = sorted(seen.values(), key=lambda x: x["score"], reverse=True)[:10]
        return {"results": final}

    except Exception as e:
        print(f"❌ Search error: {e}")
        raise HTTPException(500, str(e))
    finally:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)


# ══════════════════════════════════════════════════════════════════
# 4. CATEGORIES (Cloudinary Folders Only)
# ══════════════════════════════════════════════════════════════════
@app.post("/api/categories")
async def get_categories(user_cloudinary_url: str = Form("")):
    if not user_cloudinary_url:
        return {"categories": []}
    
    try:
        creds = get_cloudinary_creds(user_cloudinary_url)
        _configure_cloudinary(creds)
        result  = await asyncio.to_thread(cloudinary.api.root_folders)
        folders = [f["name"] for f in result.get("folders", [])]
        return {"categories": folders}
    except Exception as e:
        print(f"Category fetch error: {e}")
        return {"categories": []}


# ══════════════════════════════════════════════════════════════════
# 5. HEALTH CHECK
# ══════════════════════════════════════════════════════════════════
@app.get("/api/health")
async def health():
    return {
        "status":    "ok",
        "device":    ai.device if ai else "loading",
        "sem_slots": _inference_sem._value if _inference_sem else 0,
    }