main.py

from fastapi import FastAPI, UploadFile, File, BackgroundTasks, HTTPException, Form
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional
import uuid
import os
import shutil
import asyncio
import json
import logging
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor

from database.qdrantDB import QdrantDB
from services.embedding_service import EmbeddingService
from services.ingest_service import IngestService
from services.search_service import SearchService

app = FastAPI(title="VisionText API", version="1.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

UPLOAD_DIR = Path("uploads")
UPLOAD_DIR.mkdir(exist_ok=True)
COLLECTION_NAME = os.getenv("COLLECTION_NAME", "images")
EMBEDDING_DIM = int(os.getenv("EMBEDDING_DIM", "768"))
EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "hf-hub:timm/ViT-B-16-SigLIP")

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger(__name__)
_executor = ThreadPoolExecutor()

db = QdrantDB()
embedding_service = EmbeddingService(model_name=EMBEDDING_MODEL)
ingest_service = IngestService(db, embedding_service)
search_service = SearchService(db, embedding_service)

db.get_or_create_collection(COLLECTION_NAME, EMBEDDING_DIM)

STATIC_DIR = Path("static")
if STATIC_DIR.exists():
    app.mount("/static", StaticFiles(directory="static"), name="static")
if UPLOAD_DIR.exists():
    app.mount("/uploads", StaticFiles(directory="uploads"), name="uploads")


class SearchRequest(BaseModel):
    query: str
    top_k: int = 5


class BatchSearchRequest(BaseModel):
    queries: List[str]
    top_k: int = 5


class IngestResponse(BaseModel):
    status: str
    image_id: str
    message: str


class BatchIngestResponse(BaseModel):
    status: str
    message: str
    processed_count: int
    image_ids: List[str]


ALLOWED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".gif"}
MAX_UPLOAD_BYTES = 20 * 1024 * 1024  # 20MB


def validate_image_file(file: UploadFile):
    ext = Path(file.filename).suffix.lower()
    if ext not in ALLOWED_EXTENSIONS:
        raise HTTPException(status_code=400, detail=f"File type '{ext}' not allowed")
    if file.size and file.size > MAX_UPLOAD_BYTES:
        raise HTTPException(status_code=413, detail="File too large (max 20MB)")


def process_batch_images(image_paths: List[dict]):
    try:
        ingest_service.ingest_images_batch(COLLECTION_NAME, image_paths)
    except Exception as e:
        logger.error("Error processing batch: %s", e)


@app.get("/")
async def root():
    index_path = STATIC_DIR / "index.html"
    if index_path.exists():
        return FileResponse(index_path)
    return {
        "message": "VisionText API",
        "version": "1.0.0",
        "endpoints": {
            "search_text": "/search/text",
            "search_image": "/search/image",
            "search_batch": "/search/batch",
            "ingest_single": "/ingest/image",
            "ingest_batch": "/ingest/batch",
            "stats": "/stats",
        },
    }


@app.post("/search/text")
async def search_by_text(request: SearchRequest):
    try:
        loop = asyncio.get_event_loop()
        results = await loop.run_in_executor(
            _executor,
            lambda: search_service.search_by_text(
                collection_name=COLLECTION_NAME, query=request.query, top_k=request.top_k
            ),
        )
        return {
            "query": request.query,
            "results": [
                {"id": r.id, "score": r.score, "metadata": r.metadata} for r in results
            ],
        }
    except HTTPException:
        raise
    except Exception as e:
        logger.exception("Text search failed")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/search/image")
async def search_by_image(
    file: UploadFile = File(...),
    top_k: int = 5,
    tags: Optional[str] = Form(None),
):
    validate_image_file(file)
    temp_path = None
    tags_list = [t.strip() for t in tags.split(",") if t.strip()] if tags else None
    try:
        temp_path = UPLOAD_DIR / f"search_{uuid.uuid4()}{Path(file.filename).suffix.lower()}"
        with temp_path.open("wb") as buffer:
            shutil.copyfileobj(file.file, buffer)

        loop = asyncio.get_event_loop()
        results = await loop.run_in_executor(
            _executor,
            lambda: search_service.search_by_image(
                collection_name=COLLECTION_NAME, image_path=str(temp_path), top_k=top_k, tags=tags_list
            ),
        )
        return {
            "filename": file.filename,
            "results": [
                {"id": r.id, "score": r.score, "metadata": r.metadata} for r in results
            ],
        }
    except HTTPException:
        raise
    except Exception as e:
        logger.exception("Image search failed")
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if temp_path and temp_path.exists():
            os.remove(temp_path)


@app.post("/search/batch", response_model=dict)
async def search_batch(request: BatchSearchRequest):
    try:
        results = search_service.search_text_batch(
            COLLECTION_NAME, request.queries, request.top_k
        )
        return {
            "queries": request.queries,
            "results": [
                [{"id": r.id, "score": r.score, "metadata": r.metadata} for r in batch]
                for batch in results
            ],
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/ingest/image", response_model=IngestResponse)
async def ingest_single_image(
    file: UploadFile = File(...),
    metadata: Optional[str] = None,
):
    validate_image_file(file)
    temp_path = None
    try:
        image_id = str(uuid.uuid4())
        file_ext = Path(file.filename).suffix.lower()
        temp_path = UPLOAD_DIR / f"{image_id}{file_ext}"

        with temp_path.open("wb") as buffer:
            shutil.copyfileobj(file.file, buffer)

        meta_dict = json.loads(metadata) if metadata else {}
        meta_dict["filename"] = file.filename
        meta_dict["stored_filename"] = temp_path.name
        meta_dict["image_url"] = f"/uploads/{temp_path.name}"

        loop = asyncio.get_event_loop()
        await loop.run_in_executor(
            _executor,
            lambda: ingest_service.ingest_image(COLLECTION_NAME, image_id, str(temp_path), meta_dict),
        )

        return IngestResponse(
            status="success",
            image_id=image_id,
            message=f"Image {file.filename} ingested successfully",
        )
    except HTTPException:
        raise
    except Exception as e:
        logger.exception("Single ingest failed")
        if temp_path and temp_path.exists():
            os.remove(temp_path)
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/ingest/batch", response_model=BatchIngestResponse)
async def ingest_batch_images(
    background_tasks: BackgroundTasks,
    files: List[UploadFile] = File(...),
    default_tags: Optional[str] = Form(None),
):
    image_paths = []
    try:
        tags_list = [t.strip() for t in default_tags.split(",") if t.strip()] if default_tags else []
        image_ids = []

        for file in files:
            validate_image_file(file)
            image_id = str(uuid.uuid4())
            image_ids.append(image_id)
            file_ext = Path(file.filename).suffix.lower()
            temp_path = UPLOAD_DIR / f"{image_id}{file_ext}"

            with temp_path.open("wb") as buffer:
                shutil.copyfileobj(file.file, buffer)

            metadata = {
                "filename": file.filename,
                "stored_filename": temp_path.name,
                "image_url": f"/uploads/{temp_path.name}",
            }
            if tags_list:
                metadata["tags"] = tags_list

            image_paths.append({"id": image_id, "path": str(temp_path), "metadata": metadata})

        background_tasks.add_task(process_batch_images, image_paths)

        return BatchIngestResponse(
            status="processing",
            message="Batch ingestion started in background",
            processed_count=len(files),
            image_ids=image_ids,
        )
    except HTTPException:
        raise
    except Exception as e:
        logger.exception("Batch ingest failed")
        for item in image_paths:
            try:
                os.remove(item["path"])
            except OSError:
                pass
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/stats")
async def get_stats():
    try:
        count = db.count(COLLECTION_NAME)
        return {
            "collection_name": COLLECTION_NAME,
            "total_images": count,
            "embedding_dimension": EMBEDDING_DIM,
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.delete("/images/{image_id}")
async def delete_image(image_id: str):
    try:
        results = db.get(COLLECTION_NAME, [image_id])
        db.delete(COLLECTION_NAME, [image_id])

        if results and results[0].metadata:
            stored = results[0].metadata.get("stored_filename")
            if stored:
                file_path = UPLOAD_DIR / stored
                if file_path.exists():
                    os.remove(file_path)

        return {"status": "success", "message": f"Image {image_id} deleted"}
    except Exception as e:
        logger.exception("Delete failed for image %s", image_id)
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)