import os
import cv2
from typing import Optional
from pathlib import Path

import supervision as sv
import uvicorn
import yaml
from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.security.api_key import APIKeyHeader
from starlette.status import HTTP_403_FORBIDDEN


from groundingdino.util.inference import Model, preprocess_caption
from pdf_converter import PdfConverter


API_PARTNER_KEY = ""
API_KEY_NAME = "x-api-key"
api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)

app = FastAPI()

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

pdf_converter: PdfConverter = PdfConverter()
grounding_model = Model(
    model_config_path=os.environ.get("GROUNDING_DINO_CONFIG"),
    model_checkpoint_path=os.environ.get("GROUNDING_DINO_CHECKPOINT"),
    device="cuda:0",
)
BOX_THRESHOLD = 0.4
TEXT_THRESHOLD = 0.25

async def verify_api_key(request: Request):
    api_key = await api_key_header(request)
    if API_PARTNER_KEY is None or api_key != API_PARTNER_KEY:
        raise HTTPException(
            status_code=HTTP_403_FORBIDDEN,
            detail="Could not validate API KEY",
        )


@app.middleware("http")
async def api_key_middleware(request: Request, call_next):
    # Allow docs and openapi without API key
    if request.url.path in ["/docs", "/openapi.json", "/redoc"]:
        return await call_next(request)
    try:
        await verify_api_key(request)
    except HTTPException as exc:
        return JSONResponse(status_code=exc.status_code, content={"detail": exc.detail})
    return await call_next(request)


@app.post("/crop_ooi")
async def crop_object_of_interest(
    document_file: Optional[UploadFile] = File(
        None, description="The document to process."
    ),
    concept_list: Optional[list[str]] = Form(
        ["ID document"], description="List of concepts to detect e.g. dog, cat, rain"
    ),
    save_files: Optional[bool] = Form(
        False, description="True if crop are saved on local"
    ),
    box_threshold: Optional[float] = Form(
        0.4, description="Threshold rate to keep confidence detections"
    ),
    text_threshold: Optional[float] = Form(0.25, description="Text threshold"),
):
    content = await document_file.read()
    if document_file.content_type not in ["application/pdf", "image/jpeg", "image/png"]:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file type ({document_file.content_type}).",
        )
    try:
        images = pdf_converter.convert_pdf_bytes_to_jpg(content)
        text = "".join([f"{x}. " for x in concept_list])
        caption = preprocess_caption(text)
        for image in images:
            detections, labels = grounding_model.predict_with_caption(
                image=images,
                caption=caption,
                box_threshold=box_threshold,
                text_threshold=text_threshold,
            )
            confidences = detections.confidence.tolist()
            class_names = labels

            labels = [
                f"{class_name} {confidence:.2f}"
                for class_name, confidence in zip(class_names, confidences)
            ]
            for i, bbox in enumerate(detections.xyxy):
                x_min, y_min, x_max, y_max = tuple(bbox)
                patch = image[int(y_min) : int(y_max), int(x_min) : int(x_max)]
                # patch_img_path = str(
                #     Path("outputs") / "extract" / f"{img_path.stem}_{i:d}.png"
                # )
                # cv2.imwrite(patch_img_path, patch)

            box_annotator = sv.BoxAnnotator(color_lookup=sv.ColorLookup.INDEX)
            annotated_frame = box_annotator.annotate(
                scene=image.copy(), detections=detections
            )
            label_annotator = sv.LabelAnnotator(color_lookup=sv.ColorLookup.INDEX)
            annotated_frame = label_annotator.annotate(
                scene=annotated_frame, detections=detections, labels=labels
            )
            # cv2.imwrite(str(Path("outputs") / f"{img_path.stem}.png"), annotated_frame)

    except Exception as e:
        print(f"{e}")


if __name__ == "__main__":
    APP_PORT = int(os.environ.get("VLM_APP_PORT", 8009))
    APP_HOST = os.environ.get("VLM_APP_HOST", "0.0.0.0")
    uvicorn.run(app, host=APP_HOST, port=APP_PORT)