from dotenv import load_dotenv
import cocoindex
import datetime
import os
import requests
import base64
from fastapi import FastAPI, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from qdrant_client import QdrantClient

OLLAMA_URL = "http://localhost:11434/api/generate"
OLLAMA_MODEL = "gemma3"
QDRANT_GRPC_URL = os.getenv("QDRANT_GRPC_URL", "http://localhost:6334/")

# 1. Extract caption from image using Ollama vision model
@cocoindex.op.function(cache=True, behavior_version=1)
def get_image_caption(img_bytes: bytes) -> str:
    """
    Use Ollama's gemma3 model to extract a detailed caption from an image.
    Returns a full-sentence natural language description of the image.
    """
    img_b64 = base64.b64encode(img_bytes).decode("utf-8")
    prompt = (
        "Describe this image in one detailed, natural language sentence. "
        "Always explicitly name every visible animal species, object, and the main scene. "
        "Be specific about the type, color, and any distinguishing features. "
        "Avoid generic words like 'animal' or 'creature'—always use the most precise name (e.g., 'elephant', 'cat', 'lion', 'zebra'). "
        "If an animal is present, mention its species and what it is doing. "
        "For example: 'A large grey elephant standing in a grassy savanna, with trees in the background.'"
    )
    payload = {
        "model": OLLAMA_MODEL,
        "prompt": prompt,
        "images": [img_b64],
        "stream": False,
    }
    resp = requests.post(OLLAMA_URL, json=payload)
    resp.raise_for_status()
    result = resp.json()
    text = result.get("response", "")
    text = text.strip().replace("\n", "").rstrip(".")
    return text


# 2. Embed the caption string
@cocoindex.transform_flow()
def caption_to_embedding(caption: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[list[float]]:
    """
    Embed the caption using a CLIP model.
    This is shared logic between indexing and querying.
    """
    return caption.transform(
        cocoindex.functions.SentenceTransformerEmbed(
            model="clip-ViT-L-14",
        )
    )

# 3. CocoIndex flow: Ingest images, extract captions, embed, export to Qdrant
@cocoindex.flow_def(name="ImageObjectEmbedding")
def image_object_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope):
    data_scope["images"] = flow_builder.add_source(
        cocoindex.sources.LocalFile(path="img", included_patterns=["*.jpg", "*.jpeg", "*.png"], binary=True),
        refresh_interval=datetime.timedelta(minutes=1)  # Poll for changes every 1 minute
    )
    img_embeddings = data_scope.add_collector()
    with data_scope["images"].row() as img:
        img["caption"] = img["content"].transform(get_image_caption)
        img["embedding"] = caption_to_embedding(img["caption"])
        img_embeddings.collect(
            id=cocoindex.GeneratedField.UUID,
            filename=img["filename"],
            caption=img["caption"],
            embedding=img["embedding"],
        )
    img_embeddings.export(
        "img_embeddings",
        cocoindex.storages.Qdrant(
            collection_name="image_search",
            grpc_url=QDRANT_GRPC_URL,
        ),
        primary_key_fields=["id"],
        setup_by_user=True,
    )

# --- FastAPI app for web API ---
app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
# Serve images from the 'img' directory at /img
app.mount("/img", StaticFiles(directory="img"), name="img")

# --- CocoIndex initialization on startup ---
@app.on_event("startup")
def startup_event():
    load_dotenv()
    cocoindex.init()
    # Initialize Qdrant client
    app.state.qdrant_client = QdrantClient(
        url=QDRANT_GRPC_URL,
        prefer_grpc=True
    )
    app.state.live_updater = cocoindex.FlowLiveUpdater(image_object_embedding_flow)
    app.state.live_updater.start()

@app.get("/search")
def search(q: str = Query(..., description="Search query"), limit: int = Query(5, description="Number of results")):
    # Get the embedding for the query
    query_embedding = caption_to_embedding.eval(q)
    
    # Search in Qdrant
    search_results = app.state.qdrant_client.search(
        collection_name="image_search",
        query_vector=("embedding", query_embedding),
        limit=limit
    )
    
    # Format results
    out = []
    for result in search_results:
        out.append({
            "filename": result.payload["filename"],
            "score": result.score
        })
    return {"results": out}
