from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import os
import shutil

app = FastAPI()


skin_type_model = load_model("latest_final_skin_type_model.keras")
skin_issue_model = load_model("latest_final_skin_issues_model.keras")
skin_cancer_model = load_model("latest_final_skin_cancer_model.keras")


skin_cancer_labels = ['cancer', 'no_cancer']
skin_type_labels = ['dry', 'normal', 'oily']
skin_issue_labels = ['acne', 'no_issues', 'pigmentation', 'sensitive', 'wrinkles']

UPLOAD_FOLDER = 'uploads'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

@app.get("/")
def read_root():
    return {"message": "Skin analysis API is working ✅"}

def preprocess_image(img_path, target_size=(224, 224)):
    img = image.load_img(img_path, target_size=target_size)
    img_array = image.img_to_array(img)
    img_array = img_array / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

@app.post("/predict-skin-type")
async def predict_skin_type(image_file: UploadFile = File(...)):
    if not image_file.filename:
        raise HTTPException(status_code=400, detail="No image uploaded.")

    try:
        file_path = os.path.join(UPLOAD_FOLDER, image_file.filename)
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(image_file.file, buffer)

        img_array = preprocess_image(file_path)
        predictions = skin_type_model.predict(img_array)
        predicted_class = skin_type_labels[np.argmax(predictions)]
        confidence = float(np.max(predictions)) * 100

        result = {
            "predicted_class": predicted_class,
            "confidence": round(confidence, 2)
        }

        return JSONResponse(content=result)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if os.path.exists(file_path):
            os.remove(file_path)

@app.post("/predict-skin-issue")
async def predict_skin_issue(image_file: UploadFile = File(...)):
    if not image_file.filename:
        raise HTTPException(status_code=400, detail="No image uploaded.")

    try:
        file_path = os.path.join(UPLOAD_FOLDER, image_file.filename)
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(image_file.file, buffer)

        img_array = preprocess_image(file_path)
        predictions = skin_issue_model.predict(img_array)
        predicted_class = skin_issue_labels[np.argmax(predictions)]
        confidence = float(np.max(predictions)) * 100

        result = {
            "predicted_class": predicted_class,
            "confidence": round(confidence, 2)
        }

        return JSONResponse(content=result)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if os.path.exists(file_path):
            os.remove(file_path)

@app.post("/predict-skin-cancer")
async def predict_skin_issue(image_file: UploadFile = File(...)):
    if not image_file.filename:
        raise HTTPException(status_code=400, detail="No image uploaded.")

    try:
        file_path = os.path.join(UPLOAD_FOLDER, image_file.filename)
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(image_file.file, buffer)

        img_array = preprocess_image(file_path)
        predictions = skin_cancer_model.predict(img_array)
        predicted_class = skin_cancer_labels[np.argmax(predictions)]
        confidence = float(np.max(predictions)) * 100

        result = {
            "predicted_class": predicted_class,
            "confidence": round(confidence, 2)
        }

        return JSONResponse(content=result)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if os.path.exists(file_path):
            os.remove(file_path)


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