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:
    if api_key not in _pinecone_pool:
        if len(_pinecone_pool) >= _POOL_MAX:
            _pinecone_pool.popitem(last=False)
        _pinecone_pool[api_key] = Pinecone(api_key=api_key)
    _pinecone_pool.move_to_end(api_key)
    return _pinecone_pool[api_key]


def _configure_cloudinary(creds: dict) -> None:
    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


@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=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

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


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 (With Demo Fallback)
# ══════════════════════════════════════════════════════════════════
@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(""),
):
    # FALLBACK LOGIC: Use user keys if provided, otherwise use Space secrets
    actual_pc_key = user_pinecone_key or os.getenv("DEFAULT_PINECONE_KEY")
    actual_cld_url = user_cloudinary_url or os.getenv("DEFAULT_CLOUDINARY_URL")

    if not actual_pc_key or not actual_cld_url:
        raise HTTPException(status_code=400, detail="Cloudinary URL and Pinecone API Key are required.")

    folder    = standardize_category_name(folder_name)
    uploaded_urls = []

    cld_creds = get_cloudinary_creds(actual_cld_url)
    _configure_cloudinary(cld_creds)
    pc        = _get_pinecone(actual_pc_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)

            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}")
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)

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


# ══════════════════════════════════════════════════════════════════
# 3. SEARCH (With Demo Fallback)
# ══════════════════════════════════════════════════════════════════
@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(""),
):
    actual_pc_key = user_pinecone_key or os.getenv("DEFAULT_PINECONE_KEY")
    
    if not actual_pc_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(actual_pc_key)
        idx_obj  = pc.Index("lens-objects")
        idx_face = pc.Index("lens-faces")

        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]

        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 (With Demo Fallback)
# ══════════════════════════════════════════════════════════════════
@app.post("/api/categories")
async def get_categories(user_cloudinary_url: str = Form("")):
    actual_cld_url = user_cloudinary_url or os.getenv("DEFAULT_CLOUDINARY_URL")

    if not actual_cld_url:
        return {"categories": []}
    
    try:
        creds = get_cloudinary_creds(actual_cld_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,
    }