"""Various utilities (logger, time benchmark, args dump, numerical and stats info)"""

from pathlib import Path
from typing import Any

import numpy as np
from numpy import ndarray

from samgis_core.utilities.serialize import serialize


def _prepare_base64_input(sb: str | bytes) -> bytes:
    if isinstance(sb, str):
        # If there's any unicode here, an exception will be thrown and the function will return false
        return bytes(sb, "ascii")
    elif isinstance(sb, bytes):
        return sb
    raise ValueError("Argument must be string or bytes")


def _is_base64(sb: str | bytes) -> bool:
    import base64

    try:
        sb_bytes = _prepare_base64_input(sb)
        decoded = base64.b64decode(sb_bytes, validate=True)
        return base64.b64encode(decoded).decode("utf-8") == sb_bytes.decode("utf-8")
    except ValueError:
        return False


def base64_decode(s: str | bytes) -> str | bytes:
    """
    Decode base64 strings

    Args:
        s: input string

    Returns:
        decoded string
    """
    import base64

    if isinstance(s, str) and _is_base64(s):
        return base64.b64decode(s, validate=True).decode("utf-8")

    return s


def base64_encode(sb: str | bytes) -> bytes:
    """
    Encode input strings or bytes as base64

    Args:
        sb: input string or bytes

    Returns:
        base64 encoded bytes
    """
    import base64

    sb_bytes = _prepare_base64_input(sb)
    return base64.b64encode(sb_bytes)


def hash_calculate(
    arr_or_path: ndarray | dict[str, Any] | str | bytes | Path,
    is_file: bool,
    read_mode: str = "rb",
) -> bytes:
    """
    Return computed hash from input variable (typically a numpy array).

    Args:
        arr_or_path: input variable
        is_file: whether input is a file or not
        read_mode: mode for reading file

    Returns:
        computed hash from input variable
    """
    from hashlib import sha256
    from base64 import b64encode
    from numpy import ndarray as np_ndarray

    if is_file:
        if not isinstance(arr_or_path, (str, Path)):
            raise TypeError(
                f"arr_or_path must be str or Path when is_file=True, got {type(arr_or_path)}"
            )
        resolved = Path(arr_or_path).resolve()
        allowed_root = Path(__file__).resolve().parent.parent
        if not resolved.is_relative_to(allowed_root):
            raise ValueError(
                f"Path '{resolved}' is outside the allowed root '{allowed_root}'."
            )
        with open(resolved, read_mode) as file_to_check:
            arr_or_path = file_to_check.read()

    if isinstance(arr_or_path, np_ndarray):
        hash_fn = sha256(arr_or_path.data)
    elif isinstance(arr_or_path, dict):
        import json

        serialized = serialize(arr_or_path)
        variable_to_hash = json.dumps(serialized, sort_keys=True).encode("utf-8")
        hash_fn = sha256(variable_to_hash)
    elif isinstance(arr_or_path, str):
        hash_fn = sha256(arr_or_path.encode("utf-8"))
    elif isinstance(arr_or_path, bytes):
        hash_fn = sha256(arr_or_path)
    else:
        raise ValueError(
            f"variable 'arr':{arr_or_path} of type '{type(arr_or_path)}' not yet handled."
        )
    return b64encode(hash_fn.digest())


def normalize_array(
    arr: ndarray, new_h: int | float = 255.0, type_normalization: str = "int"
) -> ndarray:
    """
    Normalize numpy array between 0 and 'new_h' value. Default dtype of output array is int


    Args:
        arr: input numpy array
        new_h: max value of output array
        type_normalization: default dtype of output array

    Returns:
        numpy array
    """
    arr = arr.astype(float)
    arr_max = np.nanmax(arr)
    arr_min = np.nanmin(arr)
    if arr_max == arr_min:
        return np.zeros_like(arr, dtype=int if type_normalization == "int" else float)
    scaled_arr = (arr - arr_min) / (arr_max - arr_min)
    multiplied_arr = scaled_arr * new_h
    return multiplied_arr.astype(int) if type_normalization == "int" else multiplied_arr
