import os, io, sys, base64, traceback
from pathlib import Path
from flask import Flask, request, render_template, jsonify
from PIL import Image


BASE_DIR = Path(__file__).resolve().parent
if str(BASE_DIR) not in sys.path:
    sys.path.insert(0, str(BASE_DIR))

from predict import predict_pytorch, predict_tensorflow

app = Flask(__name__, template_folder=str(BASE_DIR / "templates"))
app.config["MAX_CONTENT_LENGTH"] = 5 * 1024 * 1024

ALLOWED_EXT        = {"png", "jpg", "jpeg", "webp", "bmp"}
PYTORCH_MODEL_PATH = os.getenv("PYTORCH_MODEL_PATH", str(BASE_DIR / "sara_model.pth"))
TF_MODEL_PATH      = os.getenv("TF_MODEL_PATH",      str(BASE_DIR / "sara_model.keras"))

CLASS_ICONS = {
    "buildings": "🏙️", "forest": "🌲", "glacier": "🧊",
    "mountain":  "🏔️", "sea":    "🌊", "street":  "🛣️",
}

def allowed_file(filename):
    return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXT

def image_to_b64(img):
    buf = io.BytesIO()
    img.convert("RGB").save(buf, format="JPEG")
    return base64.b64encode(buf.getvalue()).decode("utf-8")

@app.route("/", methods=["GET"])
def index():
    return render_template("index.html")

@app.route("/predict", methods=["POST"])
def predict():
    model_choice = request.form.get("model", "pytorch")
    file         = request.files.get("image")

    if not file or file.filename == "":
        return render_template("index.html", error="Upload an image."), 400
    if not allowed_file(file.filename):
        return render_template("index.html", error="You have to use JPG, PNG, WEBP ou BMP."), 400

    img_bytes = file.read()
    pil_img   = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    tmp_path  = BASE_DIR / "tmp_upload.jpg"
    pil_img.save(str(tmp_path), format="JPEG")

    try:
        if model_choice == "pytorch":
            if not Path(PYTORCH_MODEL_PATH).exists():
                raise FileNotFoundError(f"Model PyTorch not found : {PYTORCH_MODEL_PATH}")
            result = predict_pytorch(str(tmp_path), model_path=PYTORCH_MODEL_PATH)
        else:
            if not Path(TF_MODEL_PATH).exists():
                raise FileNotFoundError(f"Model TensorFlow not found : {TF_MODEL_PATH}")
            result = predict_tensorflow(str(tmp_path), model_path=TF_MODEL_PATH)

    except FileNotFoundError as e:
        tmp_path.unlink(missing_ok=True)
        return render_template("index.html", error=f"Model not found : {e}"), 500
    except Exception as e:
        tmp_path.unlink(missing_ok=True)
        print("ERREUR INFERENCE :", traceback.format_exc())
        return render_template("index.html", error=f"Error : {str(e)}"), 500
    finally:
        tmp_path.unlink(missing_ok=True)

    img_b64      = image_to_b64(pil_img)
    sorted_probs = sorted(result["all_probabilities"].items(), key=lambda x: -x[1])

    return render_template("index.html",
        result=result, model_used=model_choice,
        img_b64=img_b64, sorted_probs=sorted_probs, class_icons=CLASS_ICONS)

@app.route("/health")
def health():
    return jsonify({"status": "ok", "pytorch": Path(PYTORCH_MODEL_PATH).exists(),
                    "tensorflow": Path(TF_MODEL_PATH).exists()}), 200

if __name__ == "__main__":
    port = int(os.getenv("PORT", 7860))
    print(f"BASE_DIR : {BASE_DIR}")
    print(f"PyTorch  : {PYTORCH_MODEL_PATH} — existe : {Path(PYTORCH_MODEL_PATH).exists()}")
    print(f"TF       : {TF_MODEL_PATH} — existe : {Path(TF_MODEL_PATH).exists()}")
    app.run(host="0.0.0.0", port=port, debug=False)