#!/usr/bin/env pytest
# -*- coding: utf-8 -*-
###############################################################################
#
# Project:  GDAL/OGR Test Suite
# Purpose:  Test AVIF driver
# Author:   Even Rouault, <even dot rouault at spatialys.com>
#
###############################################################################
# Copyright (c) 2024, Even Rouault <even dot rouault at spatialys.com>
#
# SPDX-License-Identifier: MIT
###############################################################################

import base64
import os
import shutil

import gdaltest
import pytest

from osgeo import gdal

pytestmark = pytest.mark.require_driver("AVIF")


def has_avif_encoder():
    drv = gdal.GetDriverByName("AVIF")
    return drv is not None and drv.GetMetadataItem("DMD_CREATIONOPTIONLIST") is not None


def avif_encoder_supports_16_bit():
    return (
        has_avif_encoder()
        and "Valid values are 8, 10, 12 and 16"
        in gdal.GetDriverByName("AVIF").GetMetadataItem(gdal.DMD_CREATIONOPTIONLIST)
    )


def _has_geoheif_support():
    drv = gdal.GetDriverByName("AVIF")
    return drv and drv.GetMetadataItem("SUPPORTS_GEOHEIF", "AVIF")


def test_avif_subdatasets(tmp_path):

    filename = str(tmp_path / "out.avif")
    shutil.copy("data/avif/colors-animated-8bpc-alpha-exif-xmp.avif", filename)

    ds = gdal.Open(filename)
    assert ds
    assert len(ds.GetSubDatasets()) == 5
    subds1_name = ds.GetSubDatasets()[0][0]
    subds2_name = ds.GetSubDatasets()[1][0]

    ds = gdal.Open(subds1_name)
    assert ds
    assert ds.RasterXSize == 150
    assert ds.GetRasterBand(1).GetMetadataItem("STATISTICS_MINIMUM") is None
    assert ds.GetRasterBand(1).ComputeStatistics(False)
    assert ds.GetRasterBand(1).GetMetadataItem("STATISTICS_MINIMUM") is not None
    ds.Close()

    ds = gdal.Open(subds1_name)
    assert ds.GetRasterBand(1).GetMetadataItem("STATISTICS_MINIMUM") is not None

    ds = gdal.Open(subds2_name)
    assert ds
    assert ds.RasterXSize == 150
    assert ds.GetRasterBand(1).GetMetadataItem("STATISTICS_MINIMUM") is None

    with pytest.raises(Exception):
        gdal.Open(f"AVIF:0:{filename}")
    with pytest.raises(Exception):
        gdal.Open(f"AVIF:6:{filename}")
    with pytest.raises(Exception):
        gdal.Open("AVIF:1:non_existing.heic")
    with pytest.raises(Exception):
        gdal.Open("AVIF:")
    with pytest.raises(Exception):
        gdal.Open("AVIF:1")
    with pytest.raises(Exception):
        gdal.Open("AVIF:1:")


@pytest.mark.skipif(not has_avif_encoder(), reason="libavif encoder missing")
def test_avif_single_band():
    tst = gdaltest.GDALTest(
        "AVIF",
        "byte.tif",
        1,
        4672,
    )
    tst.testCreateCopy(vsimem=1, check_checksum_not_null=True, check_minmax=False)


@pytest.mark.skipif(not has_avif_encoder(), reason="libavif encoder missing")
@pytest.mark.require_driver("PNG")
def test_avif_gray_alpha():
    tst = gdaltest.GDALTest(
        "AVIF",
        "wms/gray+alpha.png",
        1,
        39910,
    )
    tst.testCreateCopy(vsimem=1, check_checksum_not_null=True, check_minmax=False)


@pytest.mark.skipif(not has_avif_encoder(), reason="libavif encoder missing")
def test_avif_rgb():
    tst = gdaltest.GDALTest(
        "AVIF",
        "rgbsmall.tif",
        1,
        21212,
        options=["QUALITY=100", "NUM_THREADS=1"],
    )
    tst.testCreateCopy(vsimem=1)


@pytest.mark.skipif(not has_avif_encoder(), reason="libavif encoder missing")
def test_avif_rgba():
    tst = gdaltest.GDALTest(
        "AVIF",
        "../../gcore/data/stefan_full_rgba.tif",
        1,
        12603,
        options=["QUALITY=100"],
    )
    tst.testCreateCopy(vsimem=1)


@pytest.mark.skipif(not has_avif_encoder(), reason="libavif encoder missing")
def test_avif_uint16():
    tst = gdaltest.GDALTest(
        "AVIF", "../../gcore/data/uint16.tif", 1, 4672, options=["NBITS=10"]
    )
    tst.testCreateCopy(vsimem=1, check_checksum_not_null=True, check_minmax=False)


@pytest.mark.skipif(not has_avif_encoder(), reason="libavif encoder missing")
@pytest.mark.parametrize("yuv_subsampling", ["444", "422", "420"])
def test_avif_yuv_subsampling(tmp_vsimem, yuv_subsampling):

    src_ds = gdal.Open("data/rgbsmall.tif")
    out_filename = str(tmp_vsimem / "out.avif")
    gdal.GetDriverByName("AVIF").CreateCopy(
        out_filename, src_ds, options=["YUV_SUBSAMPLING=" + yuv_subsampling]
    )
    ds = gdal.Open(out_filename)
    assert ds.GetMetadataItem("YUV_SUBSAMPLING", "IMAGE_STRUCTURE") == yuv_subsampling


@pytest.mark.skipif(not has_avif_encoder(), reason="libavif encoder missing")
def test_avif_nbits_from_src_ds(tmp_vsimem):

    src_ds = gdal.GetDriverByName("MEM").Create("", 1, 1, 1, gdal.GDT_UInt16)
    src_ds.GetRasterBand(1).SetMetadataItem("NBITS", "12", "IMAGE_STRUCTURE")
    out_filename = str(tmp_vsimem / "out.avif")
    gdal.GetDriverByName("AVIF").CreateCopy(out_filename, src_ds)
    ds = gdal.Open(out_filename)
    assert ds.GetRasterBand(1).GetMetadataItem("NBITS", "IMAGE_STRUCTURE") == "12"


@pytest.mark.skipif(not has_avif_encoder(), reason="libavif encoder missing")
def test_avif_exif_xmp(tmp_vsimem):

    src_ds = gdal.Open("data/avif/colors-animated-8bpc-alpha-exif-xmp.avif")
    out_filename = str(tmp_vsimem / "out.avif")
    gdal.GetDriverByName("AVIF").CreateCopy(out_filename, src_ds)
    if gdal.VSIStatL(out_filename + ".aux.xml"):
        gdal.Unlink(out_filename + ".aux.xml")
    ds = gdal.Open(out_filename)
    exif_mdd = ds.GetMetadata("EXIF")
    assert exif_mdd
    assert exif_mdd["EXIF_LensMake"] == "Google"
    xmp = ds.GetMetadata("xml:XMP")
    assert xmp
    assert xmp[0].startswith("<?xpacket")


@pytest.mark.skipif(not has_avif_encoder(), reason="libavif encoder missing")
def test_avif_icc_profile(tmp_vsimem):

    if "SOURCE_ICC_PROFILE" not in gdal.GetDriverByName("AVIF").GetMetadataItem(
        "DMD_CREATIONOPTIONLIST"
    ):
        pytest.skip("ICC profile setting requires libavif >= 1.0")

    f = open("data/sRGB.icc", "rb")
    data = f.read()
    icc = base64.b64encode(data).decode("ascii")
    f.close()

    src_ds = gdal.GetDriverByName("MEM").Create("", 1, 1, 3)
    src_ds.GetRasterBand(1).SetColorInterpretation(gdal.GCI_RedBand)
    src_ds.GetRasterBand(2).SetColorInterpretation(gdal.GCI_GreenBand)
    src_ds.GetRasterBand(3).SetColorInterpretation(gdal.GCI_BlueBand)
    src_ds.SetMetadataItem("SOURCE_ICC_PROFILE", icc, "COLOR_PROFILE")

    out_filename = str(tmp_vsimem / "out.avif")
    gdal.GetDriverByName("AVIF").CreateCopy(out_filename, src_ds)
    ds = gdal.Open(out_filename)
    assert ds.GetMetadataItem("SOURCE_ICC_PROFILE", "COLOR_PROFILE") == icc


@pytest.mark.skipif(not has_avif_encoder(), reason="libavif encoder missing")
def test_avif_creation_errors(tmp_vsimem):

    out_filename = str(tmp_vsimem / "out.avif")

    src_ds = gdal.GetDriverByName("MEM").Create("", 65537, 1)
    with pytest.raises(
        Exception,
        match="Too big source dataset. Maximum AVIF image dimension is 65,536 x 65,536 pixels",
    ):
        gdal.GetDriverByName("AVIF").CreateCopy(out_filename, src_ds)

    src_ds = gdal.GetDriverByName("MEM").Create("", 1, 65537)
    with pytest.raises(
        Exception,
        match="Too big source dataset. Maximum AVIF image dimension is 65,536 x 65,536 pixels",
    ):
        gdal.GetDriverByName("AVIF").CreateCopy(out_filename, src_ds)

    src_ds = gdal.GetDriverByName("MEM").Create("", 1, 1, 5)
    with pytest.raises(Exception, match="Unsupported number of bands"):
        gdal.GetDriverByName("AVIF").CreateCopy(out_filename, src_ds)

    src_ds = gdal.GetDriverByName("MEM").Create("", 1, 1, 1, gdal.GDT_Float32)
    with pytest.raises(
        Exception,
        match="Unsupported data type: only Byte or UInt16 bands are supported",
    ):
        gdal.GetDriverByName("AVIF").CreateCopy(out_filename, src_ds)

    src_ds = gdal.GetDriverByName("MEM").Create("", 1, 1)
    with pytest.raises(
        Exception, match="Invalid/inconsistent bit depth w.r.t data type"
    ):
        gdal.GetDriverByName("AVIF").CreateCopy(
            out_filename, src_ds, options=["NBITS=10"]
        )

    if not avif_encoder_supports_16_bit():
        src_ds = gdal.GetDriverByName("MEM").Create("", 1, 1, 1, gdal.GDT_UInt16)
        with pytest.raises(
            Exception, match="Invalid/inconsistent bit depth w.r.t data type"
        ):
            gdal.GetDriverByName("AVIF").CreateCopy(
                out_filename, src_ds, options=["NBITS=8"]
            )

    src_ds = gdal.GetDriverByName("MEM").Create("", 1, 1)
    src_ds.GetRasterBand(1).SetColorTable(gdal.ColorTable())
    with pytest.raises(
        Exception,
        match="Source dataset with color table unsupported. Use gdal_translate -expand rgb|rgba first",
    ):
        gdal.GetDriverByName("AVIF").CreateCopy(out_filename, src_ds)

    src_ds = gdal.GetDriverByName("MEM").Create("", 1, 1, 3)
    with pytest.raises(
        Exception, match="Only YUV_SUBSAMPLING=444 is supported for lossless encoding"
    ):
        gdal.GetDriverByName("AVIF").CreateCopy(
            out_filename, src_ds, options=["QUALITY=100", "YUV_SUBSAMPLING=422"]
        )

    src_ds = gdal.GetDriverByName("MEM").Create("", 1, 1)
    with pytest.raises(Exception, match="Cannot create file /i_do/not/exist.avif"):
        gdal.GetDriverByName("AVIF").CreateCopy("/i_do/not/exist.avif", src_ds)


@pytest.mark.skipif(
    not _has_geoheif_support(),
    reason="this libavif does not support opaque properties like geoheif",
)
def test_avif_geoheif_wkt2():
    ds = gdal.Open("data/heif/geo_small.avif")
    assert ds
    assert ds.RasterCount == 3
    assert ds.RasterXSize == 128
    assert ds.RasterYSize == 76
    assert ds.GetMetadataItem("NAME", "DESCRIPTION_en-AU") == "Copyright Statement"
    assert (
        ds.GetMetadataItem("DESCRIPTION", "DESCRIPTION_en-AU")
        == 'CCBY "Jacobs Group (Australia) Pty Ltd and Australian Capital Territory"'
    )
    assert ds.GetMetadataItem("TAGS", "DESCRIPTION_en-AU") == "copyright"
    assert ds.GetGeoTransform() is not None
    assert ds.GetGeoTransform() == pytest.approx(
        [691000.0, 0.1, 0.0, 6090000.0, 0.0, -0.1]
    )
    assert ds.GetSpatialRef() is not None
    assert ds.GetSpatialRef().GetAuthorityName() == "EPSG"
    assert ds.GetSpatialRef().GetAuthorityCode() == "28355"
    assert ds.GetGCPCount() == 1
    gcp = ds.GetGCPs()[0]
    assert (
        gcp.GCPPixel == pytest.approx(0, abs=1e-5)
        and gcp.GCPLine == pytest.approx(0, abs=1e-5)
        and gcp.GCPX == pytest.approx(691000.0, abs=1e-5)
        and gcp.GCPY == pytest.approx(6090000.0, abs=1e-5)
        and gcp.GCPZ == pytest.approx(0, abs=1e-5)
    )


@pytest.mark.skipif(
    not _has_geoheif_support(),
    reason="this libavif does not support opaque properties like geoheif",
)
def test_avif_geoheif_uri():
    ds = gdal.Open("data/heif/geo_crsu.avif")
    assert ds
    assert ds.RasterCount == 3
    assert ds.RasterXSize == 256
    assert ds.RasterYSize == 64
    assert ds.GetMetadataItem("NAME", "DESCRIPTION_en-AU") == "Copyright Statement"
    assert (
        ds.GetMetadataItem("DESCRIPTION", "DESCRIPTION_en-AU")
        == 'CCBY "Jacobs Group (Australia) Pty Ltd and Australian Capital Territory"'
    )
    assert ds.GetMetadataItem("TAGS", "DESCRIPTION_en-AU") == "copyright"
    assert ds.GetGeoTransform() is not None
    assert ds.GetGeoTransform() == pytest.approx(
        [691051.2, 0.1, 0.0, 6090000.0, 0.0, -0.1]
    )
    assert ds.GetSpatialRef() is not None
    assert ds.GetSpatialRef().GetAuthorityName() == "EPSG"
    assert ds.GetSpatialRef().GetAuthorityCode() == "32755"

    assert ds.GetGCPCount() == 1
    gcp = ds.GetGCPs()[0]
    assert (
        gcp.GCPPixel == pytest.approx(0, abs=1e-5)
        and gcp.GCPLine == pytest.approx(0, abs=1e-5)
        and gcp.GCPX == pytest.approx(691051.2, abs=1e-5)
        and gcp.GCPY == pytest.approx(6090000.0, abs=1e-5)
        and gcp.GCPZ == pytest.approx(0, abs=1e-5)
    )


@pytest.mark.skipif(
    not _has_geoheif_support(),
    reason="this libavif does not support opaque properties like geoheif",
)
def test_avif_geoheif_curie():
    ds = gdal.Open("data/heif/geo_curi.avif")
    assert ds
    assert ds.RasterCount == 3
    assert ds.RasterXSize == 256
    assert ds.RasterYSize == 64
    assert ds.GetMetadataItem("NAME", "DESCRIPTION_en-AU") == "Copyright Statement"
    assert (
        ds.GetMetadataItem("DESCRIPTION", "DESCRIPTION_en-AU")
        == 'CCBY "Jacobs Group (Australia) Pty Ltd and Australian Capital Territory"'
    )
    assert ds.GetMetadataItem("TAGS", "DESCRIPTION_en-AU") == "copyright"
    assert ds.GetSpatialRef() is not None
    assert ds.GetSpatialRef().GetAuthorityName() == "EPSG"
    assert ds.GetSpatialRef().GetAuthorityCode() == "32755"

    assert ds.GetGeoTransform() is not None
    assert ds.GetGeoTransform() == pytest.approx(
        [691051.2, 0.1, 0.0, 6090000.0, 0.0, -0.1]
    )
    assert ds.GetGCPCount() == 1
    gcp = ds.GetGCPs()[0]
    assert (
        gcp.GCPPixel == pytest.approx(0, abs=1e-5)
        and gcp.GCPLine == pytest.approx(0, abs=1e-5)
        and gcp.GCPX == pytest.approx(691051.2, abs=1e-5)
        and gcp.GCPY == pytest.approx(6090000.0, abs=1e-5)
        and gcp.GCPZ == pytest.approx(0, abs=1e-5)
    )


###############################################################################


@pytest.mark.skipif(not has_avif_encoder(), reason="libavif encoder missing")
def test_avif_close(tmp_path):

    ds = gdal.GetDriverByName("AVIF").CreateCopy(
        tmp_path / "out.avif", gdal.Open("data/rgbsmall.tif")
    )
    ds.Close()
    os.remove(tmp_path / "out.avif")


###############################################################################


@pytest.mark.skipif(
    not avif_encoder_supports_16_bit(), reason="libavif encoder does not support 16 bit"
)
@pytest.mark.parametrize("depth", [10, 12, 16])
@pytest.mark.parametrize("quality", [60, 100])
def test_avif_16bit(tmp_vsimem, depth, quality):

    max_val = (1 << depth) - 1
    src_ds = gdal.Translate(
        "",
        "data/byte.tif",
        options=f"-of MEM -scale 0 255 0 {max_val} -co NBITS={depth} -ot UInt16",
    )

    gdal.GetDriverByName("AVIF").CreateCopy(
        tmp_vsimem / "out.avif", src_ds, options=[f"QUALITY={quality}"]
    )

    gdal.Unlink(tmp_vsimem / "out.avif.aux.xml")

    ds = gdal.Open(tmp_vsimem / "out.avif")
    assert ds.GetRasterBand(1).DataType == gdal.GDT_UInt16
    assert (
        ds.GetRasterBand(1).GetMetadataItem("NBITS", "IMAGE_STRUCTURE") == None
        if depth == 16
        else str(depth)
    )
    if quality == 100:
        assert ds.GetRasterBand(1).ReadRaster() == src_ds.GetRasterBand(1).ReadRaster()
    else:
        assert ds.GetRasterBand(1).ReadRaster() != src_ds.GetRasterBand(1).ReadRaster()
        assert ds.GetRasterBand(1).ComputeRasterMinMax(False) == pytest.approx(
            src_ds.GetRasterBand(1).ComputeRasterMinMax(False), rel=0.1
        )
