from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import numpy as np
import cv2
import io
import base64
from PIL import Image
import image_processing

app = FastAPI(
    title="Soil Image Classification API",
    description="Classifies soil images based on visual features or segmented regions.",
    version="1.0"
)

# Allow CORS (helpful for frontend dev)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # You can restrict this to your frontend domain
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Convert image bytes to OpenCV format
def read_imagefile(image_bytes):
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

# Encode image (OpenCV format) to base64
def encode_image_to_base64(image):
    _, buffer = cv2.imencode('.jpg', image)
    return base64.b64encode(buffer).decode('utf-8')

@app.get("/")
def root():
    return {"message": "Soil Image Classifier API is running."}

# Whole image prediction

@app.post("/predictsoil")
@app.post("/predictsoil/")
async def predict_image(file: UploadFile = File(...)):
    try:
        image_bytes = await file.read()
        image = read_imagefile(image_bytes)

        features_df, predicted_class, confidence = image_processing.predict_image_class_with_features(
            cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        )

        return {
            "predicted_class": predicted_class,
            "confidence": float(confidence),
            "features": features_df.to_dict(orient="records")[0]
        }

    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})

# Region-based prediction
@app.post("/predict-regions")
@app.post("/predict-regions/")
async def predict_regions(file: UploadFile = File(...), k_clusters: int = 2):
    try:
        image_bytes = await file.read()
        image = read_imagefile(image_bytes)

        segmented_image, region_predictions = image_processing.segment_and_classify_regions(
            cv2.cvtColor(image, cv2.COLOR_BGR2RGB), 
            k_clusters=k_clusters
        )

        base64_segmented = encode_image_to_base64(cv2.cvtColor(segmented_image, cv2.COLOR_BGR2RGB))

        results = []
        for region in region_predictions:
            results.append({
                "class": region["class"],
                "confidence": float(region["confidence"]),
                "bbox": {
                    "x": region["bbox"][0],
                    "y": region["bbox"][1],
                    "width": region["bbox"][2],
                    "height": region["bbox"][3],
                }
            })

        return {
            "region_count": len(results),
            "regions": results,
            "segmented_image_base64": base64_segmented
        }

    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})
    

@app.get("/")
def home():
    return {"message": "FastAPI is running on Hugging Face Spaces!"}