from fastapi import FastAPI, UploadFile, status
from fastapi.responses import JSONResponse
from ultralytics import YOLO
import cv2
import numpy as np
from fastapi.middleware.cors import CORSMiddleware
import base64
import math

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["POST"],
    allow_headers=["content-type", "accept"],
)

model = YOLO("best.pt")

@app.post("/inference")
async def inference(file: UploadFile):
    if file.content_type != "image/jpeg" and file.content_type != "image/png" and file.content_type != "image/jpg":
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={"error": "Invalid file format"}
        )
    image_bytes = await file.read()
    img = np.frombuffer(image_bytes, dtype=np.uint8)
    img = cv2.imdecode(img, cv2.IMREAD_COLOR)
    results = model.predict(source=img, conf=0.3)
    detections = []
    for r in results: 
        boxes = r.boxes
    for box in boxes: 
        label = box.cls[0].item()  
        label_name = model.names[label]  
        
        # Skip NoTumor class - don't draw bounding box
        if label_name.lower() == "notumor":
            continue
        
        x1, y1, x2, y2 = box.xyxy[0]
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
    
        cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 255), 2)

        detections.append(label_name)
        (text_width, text_height), baseline = cv2.getTextSize(label_name, cv2.FONT_HERSHEY_SIMPLEX, 0.9, 2)
        cv2.rectangle(img, (x1, y1 - text_height - baseline - 5), (x1 + text_width, y1), (255, 0, 255), -1)
        cv2.putText(img, label_name, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 0), 2)
    
    resp_img_bytes = cv2.imencode('.jpg', img)[1].tobytes()
    img_base64 = base64.b64encode(resp_img_bytes).decode('utf-8')
    return JSONResponse({
        "image": img_base64,
        "detections": detections
    })


AVERAGE_TUMOR_VOLUME = 5751.46 

def calculate_sphere_volume(width, height):
    """
    Menghitung volume tumor menggunakan rumus bola
    Volume = (4/3) * π * r³
    Diameter = rata-rata dari width dan height
    """
    try:
        diameter = (width + height) / 2
        radius = diameter / 2
        volume = (4/3) * math.pi * (radius ** 3)
        result = round(volume, 2) 
        result = result * 0.00757132152125087872521790538421
        if result < 523: result = 523
        return result
    except Exception as e:
        print(f"Terjadi error: {e}")
        return None

@app.post("/inference_volume")
async def inference_volume(file: UploadFile):
    """
    Endpoint sederhana untuk deteksi tumor dengan volume
    Return: JSON dengan volume_mm3, class, dan image_bytes
    """
    if file.content_type not in ["image/jpeg", "image/png", "image/jpg"]:
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={"error": "Invalid file format"}
        )
    
    image_bytes = await file.read()
    img = np.frombuffer(image_bytes, dtype=np.uint8)
    img = cv2.imdecode(img, cv2.IMREAD_COLOR)
    
    results = model.predict(
        source=img, 
        conf=0.5,
        iou=0.2
    )
    
    detections = []
    
    for r in results: 
        boxes = r.boxes
        
        for box in boxes: 
            label = int(box.cls[0].item())
            label_name = model.names[label]
            
            # Skip NoTumor class - don't draw bounding box or include in detections
            if label_name.lower() == "notumor":
                continue
            
            x1, y1, x2, y2 = box.xyxy[0]
            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
            
            width = x2 - x1
            height = y2 - y1
            
            # Hitung volume
            volume_mm3 = calculate_sphere_volume(width, height)
            print("VOLUME")
            print(volume_mm3)
            
            detections.append({
                "class": label_name,
                "volume_mm3": volume_mm3
            })
            
            # Draw bounding box
            cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 255), 2)
            
            # Label dengan volume
            text = f"{label_name}"
            vol_text = f"{volume_mm3} mm3"
            
            (text_width, text_height), baseline = cv2.getTextSize(
                text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2
            )
            
            cv2.rectangle(
                img, 
                (x1, y1 - text_height - baseline - 25), 
                (x1 + max(text_width, 100), y1), 
                (255, 0, 255), 
                -1
            )
            
            cv2.putText(
                img, text, (x1, y1 - 20), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2
            )
            cv2.putText(
                img, vol_text, (x1, y1 - 5), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1
            )
    
    # Encode image ke bytes
    _, buffer = cv2.imencode('.jpg', img)
    img_bytes = buffer.tobytes()
    
    # Convert ke base64 untuk JSON
    img_base64 = base64.b64encode(img_bytes).decode('utf-8')
    
    response = {
        "detections": detections,
        "image_bytes": img_base64
    }
    
    return JSONResponse(content=response)