from fastapi import FastAPI, File, UploadFile, HTTPException, Query
from fastapi.responses import JSONResponse, StreamingResponse, RedirectResponse
from ultralytics import YOLO
import cv2
import numpy as np
import asyncio
from concurrent.futures import ThreadPoolExecutor
import torch
from PIL import Image
import io
import hashlib
from enum import Enum

app = FastAPI(
    title="Door & Window Detection API",
    description="Upload an image (.jpg or .png) to detect doors and windows. Select response type: JSON or image.",
    version="1.0.0"
)

model = None
executor = None

class ResponseType(str, Enum):
    empty = ""
    json = "json"
    image = "image"

def load_model():
    global model
    model = YOLO("yolov8m_custom.pt")
    model.to('cuda' if torch.cuda.is_available() else 'cpu')
    dummy_image = np.zeros((640, 640, 3), dtype=np.uint8)
    model.predict(dummy_image, verbose=False, conf=0.5, iou=0.45)

@app.on_event("startup")
async def startup_event():
    global executor
    executor = ThreadPoolExecutor(max_workers=2)
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(executor, load_model)

@app.on_event("shutdown")
async def shutdown_event():
    global executor
    if executor:
        executor.shutdown(wait=True)

def generate_color_for_label(label: str) -> tuple:
    hash_value = int(hashlib.md5(label.encode()).hexdigest(), 16)
    np.random.seed(hash_value % (2**32))
    color = tuple(int(c) for c in np.random.randint(0, 256, 3))
    return color

def process_image_sync(image_bytes: bytes) -> dict:
    pil_image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
    image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
    height, width = image.shape[:2]
    max_size = 1280

    if max(height, width) > max_size:
        scale = max_size / float(max(height, width))
        image = cv2.resize(image, (int(width * scale), int(height * scale)))

    results = model.predict(
        image, 
        verbose=False,
        conf=0.5,
        iou=0.45,
        max_det=100,
        half=torch.cuda.is_available()
    )

    detections = []
    result = results[0]
    if result.boxes:
        for i in range(len(result.boxes)):
            label = model.names[int(result.boxes.cls[i])]
            confidence = float(result.boxes.conf[i])
            if confidence >= 0.5:
                x1, y1, x2, y2 = result.boxes.xyxy[i].tolist()
                bbox = [x1, y1, x2 - x1, y2 - y1]
                pt1, pt2 = (int(x1), int(y1)), (int(x2), int(y2))
                color = generate_color_for_label(label)

                cv2.rectangle(image, pt1, pt2, color, 2)
                label_text = f"{label}: {confidence:.2f}"
                text_size = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)[0]
                cv2.rectangle(image, (int(x1), int(y1) - text_size[1] - 4),
                              (int(x1) + text_size[0], int(y1)), color, -1)
                cv2.putText(image, label_text, (int(x1), int(y1) - 2),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

                detections.append({
                    "label": label,
                    "confidence": round(confidence, 2),
                    "bbox": [round(coord, 2) for coord in bbox]
                })

    pil_result_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    img_byte_arr = io.BytesIO()
    pil_result_image.save(img_byte_arr, format='PNG')
    img_byte_arr.seek(0)

    return {"detections": detections, "image": img_byte_arr.getvalue()}

@app.get("/")
async def redirect_to_docs():
    return RedirectResponse(url="/docs")

@app.post("/predict")
async def predict(
    file: UploadFile = File(..., description="Upload a JPG or PNG image"),
    response_type: ResponseType = Query(..., description="Select response type: 'json' or 'image'")
):
    content = await file.read()
    if len(content) == 0:
        raise HTTPException(status_code=400, detail="Empty file uploaded")
    if len(content) > 10 * 1024 * 1024:
        raise HTTPException(status_code=400, detail="File size too large. Maximum 10MB allowed.")

    loop = asyncio.get_event_loop()
    result = await loop.run_in_executor(executor, process_image_sync, content)

    if response_type == ResponseType.image:
        return StreamingResponse(
            io.BytesIO(result["image"]),
            media_type="image/png",
            headers={
                "Content-Disposition": "inline; filename=detected_objects.png",
                "X-Detection-Count": str(len(result["detections"]))
            }
        )
    else:
        return JSONResponse(
            content={"detections": result["detections"]},
            headers={"X-Detection-Count": str(len(result["detections"]))}
        )