import io
import time

import requests
import streamlit as st
from openfoodfacts.images import generate_image_url
from PIL import Image


@st.cache_data
def send_prediction_request(image_url: str, model_name: str, server_base_url: str):
    return requests.get(
        f"{server_base_url}/api/v1/images/predict",
        params={"image_url": image_url, "models": model_name, "output_image": 1},
    )


def get_product(barcode: str):
    r = requests.get(f"https://world.openfoodfacts.org/api/v2/product/{barcode}")

    if r.status_code == 404:
        return None

    return r.json()["product"]


def run(barcode: str, model_names: list[str], server_base_url: str):
    product = get_product(barcode)
    st.markdown(f"[Product page](https://world.openfoodfacts.org/product/{barcode})")

    if not product:
        st.error(f"Product {barcode} not found")
        return

    images = product.get("images", [])

    if not images:
        st.error(f"No images found for product {barcode}")
        return

    for image_id, _ in images.items():
        if not image_id.isdigit():
            continue

        image_url = generate_image_url(barcode, f"{image_id}")

        for model_name in model_names:
            start = time.monotonic()
            response = send_prediction_request(image_url, model_name, server_base_url)
            elapsed = time.monotonic() - start

            if response.headers["Content-Type"] != "image/jpeg":
                st.error(f"Error: {response.text}")
                continue
            image = Image.open(io.BytesIO(response.content))
            st.write(f"Image {image_id}")
            st.image(image, caption=f"Model: {model_name} ({elapsed:.2f}s)")
        st.divider()


st.title("Object detection demo")
st.markdown(
    "This Streamlit is useful to test object detection models running in production at Open Food Facts."
)
default_barcode = st.query_params["barcode"] if "barcode" in st.query_params else ""
model_names = st.multiselect(
    "Models",
    options=[
        "nutrition-table-yolo",
        "nutrition-table",
        "nutriscore",
        "nutriscore-yolo",
        "universal-logo-detector",
    ],
    help="Select the model(s) to use",
    default=["nutrition-table-yolo", "nutrition-table"],
)
barcode = st.text_input(
    "barcode", help="Barcode of the product", value=default_barcode
).strip()
st.query_params["barcode"] = barcode

# Default server is staging
server_base_url = "https://robotoff.openfoodfacts.net"

if "env" in st.query_params:
    if st.query_params["env"] == "prod":
        server_base_url = "https://robotoff.openfoodfacts.net"
    elif st.query_params["env"] == "dev":
        server_base_url = "http://localhost:5000"


if barcode:
    run(barcode, model_names, server_base_url)