"""functions using machine learning instance model(s)"""

from datetime import datetime
from os import getenv
from pathlib import Path
from typing import Any

from samgis_core import app_logger
from samgis_core.prediction_api.ports import PredictorPort
from samgis_core.prediction_api.sam2_adapter import Sam2OnnxPredictor
from samgis_core.utilities.type_hints import ListDict
from samgis_web import MODEL_FOLDER
from samgis_web.io_package import raster_helpers
from samgis_web.io_package.geo_helpers import get_vectorized_raster_as_geojson
from samgis_web.io_package.tms2geotiff import download_extent
from samgis_web.utilities.constants import (
    DEFAULT_INPUT_WIDTH,
    DEFAULT_URL_TILES,
    MODEL_NAME,
    SLOPE_CELLSIZE,
)
from samgis_web.web.web_helpers import check_source_type_is_terrain

type LlistFloat = list[list[float]]
type DictStrInt = dict[str, str | int]

__all__ = [
    "samexporter_predict",
    "models_dict",
]

models_dict: dict[str, dict[str, PredictorPort | None]] = {}


def samexporter_predict(
    bbox: LlistFloat,
    prompt: ListDict,
    zoom: float,
    model_name: str = MODEL_NAME,
    source: Any = DEFAULT_URL_TILES,
    source_name: str | None = None,
    model_folder: str | Path = MODEL_FOLDER,
) -> DictStrInt:
    """
    Return predictions as a geojson from a geo-referenced image using the given input prompt.

    1. if necessary instantiate a segment anything machine learning instance model
    2. download a geo-referenced raster image delimited by the coordinates bounding box (bbox)
    3. get a prediction image from the segment anything instance model using the input prompt
    4. get a geo-referenced geojson from the prediction image

    Args:
        bbox: coordinates bounding box
        prompt: machine learning input prompt
        zoom: Level of detail
        model_name: machine learning model name
        source: xyz tile provider object
        source_name: name of tile provider,
        model_folder: ML models folder

    Returns:
        dict containing the output geojson, the prediction masks number and the geojson shapes number

    """
    import numpy as np

    if model_name not in models_dict or models_dict[model_name]["instance"] is None:
        app_logger.info(f"missing instance model {model_name}, instantiating it now!")
        model_instance = Sam2OnnxPredictor(model_dir=model_folder)
        models_dict[model_name] = {"instance": model_instance}
    app_logger.debug(f"using a {model_name} instance model...")
    models_instance = models_dict[model_name]["instance"]
    if models_instance is None:
        raise RuntimeError(f"failed to instantiate model '{model_name}'")

    pt0, pt1 = bbox
    folder_write_tmp_on_disk = getenv("WRITE_TMP_ON_DISK", "")
    app_logger.info(f"folder_write_tmp_on_disk:{folder_write_tmp_on_disk}.")
    now = datetime.now().strftime("%Y%m%d_%H%M%S")
    prefix = f"w{pt1[1]},s{pt1[0]},e{pt0[1]},n{pt0[0]}_"
    app_logger.info(
        f"tile_source: {source}: downloading geo-referenced raster with bbox {bbox}, zoom {zoom}."
    )
    img, transform = download_extent(
        w=pt1[1], s=pt1[0], e=pt0[1], n=pt0[0], zoom=int(zoom), source=source
    )
    if bool(folder_write_tmp_on_disk):
        if not (img.shape and len(img.shape) == 3 and img.shape[2] == 3):
            raise ValueError(f"wrong image shape: '{img.shape}'")
        raster_helpers.write_raster_png(
            img,
            transform,
            f"{source_name}_{prefix}_{now}_",
            "raw1",
            folder_write_tmp_on_disk,
        )

    if check_source_type_is_terrain(source):
        app_logger.info("terrain-rgb like raster: transforms it into a DEM")
        dem = raster_helpers.get_raster_terrain_rgb_like(img, source.name)
        # set a slope cell size proportional to the image width
        slope_cellsize = int(img.shape[1] * SLOPE_CELLSIZE / DEFAULT_INPUT_WIDTH)
        app_logger.info(
            f"terrain-rgb like raster: compute slope, curvature using {slope_cellsize} as cell size."
        )
        img = raster_helpers.get_rgb_prediction_image(dem, slope_cellsize)
        if bool(folder_write_tmp_on_disk):
            if not (img.shape and len(img.shape) == 3 and img.shape[2] == 3):
                raise ValueError(f"wrong image shape: '{img.shape}'")
            raster_helpers.write_raster_png(
                img,
                transform,
                f"{source_name}_{prefix}_{now}_",
                "rgb2",
                folder_write_tmp_on_disk,
            )
        if not (dem.shape and len(dem.shape) == 2):
            raise ValueError(f"wrong img (DEM) shape: '{dem.shape}'")
        if bool(folder_write_tmp_on_disk):
            dem = np.nan_to_num(dem, nan=0).astype(np.int16)
            raster_helpers.write_raster_tiff(
                dem,
                transform,
                f"{source_name}_{prefix}_{now}_",
                "raw3",
                folder_write_tmp_on_disk,
            )
    app_logger.info(
        f"img type {type(img)} with shape/size:{img.size}, transform type: {type(transform)}, transform:{transform}."
    )
    app_logger.info(f"source_name:{source_name}, source_name type:{type(source_name)}.")
    models_instance.set_image(img)
    masks, ious = models_instance.predict(prompt)
    best = int(np.argmax(ious))
    mask = masks[best]
    n_predictions = len(masks)

    if bool(folder_write_tmp_on_disk):
        from PIL.Image import fromarray as pil_fromarray

        mask_pil = pil_fromarray(mask)
        mask_pil.save(
            Path(folder_write_tmp_on_disk)
            / f"{source_name}_{prefix}_{now}__mask_row.png"
        )

    app_logger.info(
        f"created {n_predictions} masks, type {type(mask)}, size {mask.size}: preparing geojson conversion"
    )
    app_logger.info(f"mask shape:{mask.shape}.")
    geojson_content = get_vectorized_raster_as_geojson(mask, transform)
    if bool(folder_write_tmp_on_disk):
        geojson = str(geojson_content["geojson"])
        raster_helpers.write_geojson_on_disk(
            geojson,
            f"{source_name}_{prefix}_{now}_",
            "geojson",
            folder_write_tmp_on_disk,
        )
    return {"n_predictions": n_predictions, **geojson_content}
