from fastapi import FastAPI, File, UploadFile, Form
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
import cv2
import numpy as np
import base64
import time
import os

app = FastAPI(title="Construction Detection API")

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

# Load YOLO model
HF_REPO_ID = "newtechdevng/construction_detection_fine_tune"
MODEL_FILE  = "best_v2_finetune.pt"
model_path  = hf_hub_download(repo_id=HF_REPO_ID, filename=MODEL_FILE)
model       = YOLO(model_path)

# ArUco setup
ARUCO_DICT     = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50)
ARUCO_PARAMS   = cv2.aruco.DetectorParameters()
ARUCO_DETECTOR = cv2.aruco.ArucoDetector(ARUCO_DICT, ARUCO_PARAMS)

CLASS_COLORS = {
    "beam":    (255, 100,   0),
    "column":  (  0, 255, 255),
    "door":    (255,   0, 255),
    "floor":   (  0, 255,   0),
    "stairs":  (255, 255,   0),
    "wall":    (  0, 100, 255),
    "window":  (100,   0, 255),
}

def detect_aruco_scale(img, marker_size_cm=10.0):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    corners, ids, _ = ARUCO_DETECTOR.detectMarkers(gray)
    if ids is None:
        return None, None
    marker_corners = corners[0][0]
    w_px = np.linalg.norm(marker_corners[0] - marker_corners[1])
    h_px = np.linalg.norm(marker_corners[1] - marker_corners[2])
    pixels_per_cm = (w_px + h_px) / 2 / marker_size_cm
    return float(pixels_per_cm), corners

@app.get("/")
def root():
    return {
        "model": MODEL_FILE,
        "classes": list(CLASS_COLORS.keys()),
        "calibration": "Auto via ArUco marker on hard hat (10cm x 10cm)",
        "endpoints": {
            "POST /detect": "Send image → get detections in cm",
            "GET  /health": "Health check"
        }
    }

@app.get("/health")
def health():
    return {"status": "ok", "model": MODEL_FILE}

@app.post("/detect")
async def detect(
    file: UploadFile = File(...),
    marker_size_cm: float = Form(10.0),
    confidence: float = Form(0.2),
    iou: float = Form(0.3)
):
    start = time.time()

    contents = await file.read()
    nparr    = np.frombuffer(contents, np.uint8)
    img      = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    # ArUco auto-calibration
    pixels_per_cm, aruco_corners = detect_aruco_scale(img, marker_size_cm)
    calibrated = pixels_per_cm is not None

    # Draw ArUco marker highlight
    if calibrated:
        cv2.aruco.drawDetectedMarkers(img, aruco_corners)

    # Run YOLO with lower confidence + iou for more detections
    results    = model(img, conf=confidence, iou=iou)[0]
    detections = []

    for box in results.boxes:
        x1, y1, x2, y2 = map(int, box.xyxy[0])
        cls   = results.names[int(box.cls[0])]
        conf  = round(float(box.conf[0]), 2)
        w_px  = x2 - x1
        h_px  = y2 - y1
        color = CLASS_COLORS.get(cls, (0, 255, 0))

        w_cm = round(float(w_px) / pixels_per_cm, 1) if calibrated else None
        h_cm = round(float(h_px) / pixels_per_cm, 1) if calibrated else None

        # Draw bounding box
        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)

        # Label
        label = f"{cls} {conf:.2f}"
        if calibrated:
            label += f" | {w_cm}x{h_cm}cm"

        # Background for label text
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 2)
        cv2.rectangle(img, (x1, y1 - th - 10), (x1 + tw, y1), color, -1)
        cv2.putText(img, label, (x1, y1 - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 0), 2)

        detections.append({
            "class":      cls,
            "confidence": conf,
            "bbox":       [int(x1), int(y1), int(x2), int(y2)],
            "width_px":   int(w_px),
            "height_px":  int(h_px),
            "width_cm":   round(float(w_cm), 1) if w_cm is not None else None,
            "height_cm":  round(float(h_cm), 1) if h_cm is not None else None,
        })

    # Encode result image
    _, buf  = cv2.imencode(".jpg", img)
    img_b64 = base64.b64encode(buf).decode()

    return {
        "success":          True,
        "calibrated":       bool(calibrated),
        "pixels_per_cm":    round(pixels_per_cm, 2) if calibrated else None,
        "marker_size_cm":   float(marker_size_cm),
        "inference_time_s": round(float(time.time() - start), 3),
        "total":            int(len(detections)),
        "detections":       detections,
        "image_base64":     img_b64,
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)