from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse, JSONResponse
import os
import uvicorn

app = FastAPI()

MODEL_DIRECTORY = "models"

# Code-to-plant mapping starting from "000001"
CODE_TO_PLANT = {
    "000001": "Apple",
    "000002": "Cherry",
    "000003": "Corn",
    "000004": "Grape",
    "000005": "Peach",
    "000006": "Pepperbell",
    "000007": "Potato",
    "000008": "Rice",
    "000009": "Strawberry",
    "000010": "Tomato"
}

@app.get("/health")
async def api_health_check():
    return JSONResponse(content={"status": "Service is running"})
    
@app.get("/download/{plant_name}")
async def download_model(plant_name: str):
    filename = f"model_{plant_name}.tflite"
    file_path = os.path.join(MODEL_DIRECTORY, filename)

    # Check if file exists
    if not os.path.isfile(file_path):
        raise HTTPException(status_code=404, detail=f"Model file '{filename}' not found")

    return FileResponse(file_path, filename=filename, media_type="application/octet-stream")

@app.get("/download_by_code/{code}")
async def download_model_by_code(code: str):
    plant_name = CODE_TO_PLANT.get(code)
    if not plant_name:
        raise HTTPException(status_code=404, detail="Invalid 6-digit code provided")

    filename = f"model_{plant_name}.tflite"
    file_path = os.path.join(MODEL_DIRECTORY, filename)

    if not os.path.isfile(file_path):
        raise HTTPException(status_code=404, detail=f"Model file '{filename}' not found")

    return FileResponse(file_path, filename=filename, media_type="application/octet-stream")

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