"""cog_fixtures.cog module"""
import abc
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Optional, Sequence, Union

import affine
import rasterio
from fake_geo_images.fakegeoimages import FakeGeoImage
from rasterio.crs import CRS
from rasterio.io import MemoryFile
from rio_cogeo.cogeo import cog_translate


@dataclass  # type: ignore
class ImageBase(abc.ABC):
    """base image class"""

    def __enter__(self):
        """support context management"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """support context management"""
        self.close()

    @abc.abstractmethod
    def __post_init__(self):
        """post init hook"""
        ...

    @abc.abstractmethod
    def close(self):
        """close resources"""
        ...


@dataclass
class FakeImage(ImageBase):
    """fake rasterio image"""

    driver: str
    width: int
    height: int
    count: int
    dtype: str
    crs: CRS
    transform: affine.Affine
    nodata: Union[float, int] = None

    def __post_init__(self):
        """post init hook"""
        fpath, _ = FakeGeoImage(
            xsize=self.width,
            ysize=self.height,
            num_bands=self.count,
            data_type=self.dtype,
            out_dir=Path("/vsimem"),
            crs=self.crs.to_epsg(),
            nodata=self.nodata or -1,
        ).create(transform=self.transform)
        self.handle = rasterio.open(fpath)

    def close(self):
        """close resources"""
        self.handle.close()


@dataclass
class FakeCog(ImageBase):
    """fake cloud optimized geotiff"""

    src_img: FakeImage
    dst_kwargs: Dict
    indexes: Optional[Sequence[int]] = None
    nodata: Optional[Union[int, float]] = None
    dtype: Optional[str] = None
    add_mask: bool = False
    overview_level: Optional[int] = None
    overview_resampling: str = "nearest"
    web_optimized: bool = False
    latitude_adjustment: bool = True
    resampling: str = "nearest"
    config: Optional[Dict] = None
    allow_intermediate_compression: bool = False
    forward_band_tags: bool = False
    quiet: bool = False
    temporary_compression: str = "DEFLATE"

    def __post_init__(self):
        """post init hook"""
        self.memfile = MemoryFile()
        cog_translate(
            source=self.src_img.handle,
            dst_path=self.memfile.name,
            dst_kwargs=self.dst_kwargs,
            indexes=self.indexes,
            nodata=self.nodata,
            dtype=self.dtype,
            add_mask=self.add_mask,
            overview_level=self.overview_level,
            overview_resampling=self.overview_resampling,
            web_optimized=self.web_optimized,
            latitude_adjustment=self.latitude_adjustment,
            resampling=self.resampling,
            in_memory=True,
            config=self.config,
            allow_intermediate_compression=self.allow_intermediate_compression,
            forward_band_tags=self.forward_band_tags,
            quiet=self.quiet,
            temporary_compression=self.temporary_compression,
        )
        self.handle = self.memfile.open()

    def close(self):
        """close resources"""
        self.handle.close()
        self.memfile.close()

"""cog_fixtures.cog module"""
import abc
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Optional, Sequence, Union

import affine
import rasterio
from fake_geo_images.fakegeoimages import FakeGeoImage
from rasterio.crs import CRS
from rasterio.io import MemoryFile
from rio_cogeo.cogeo import cog_translate


@dataclass  # type: ignore
class ImageBase(abc.ABC):
    """base image class"""

    def __enter__(self):
        """support context management"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """support context management"""
        self.close()

    @abc.abstractmethod
    def __post_init__(self):
        """post init hook"""
        ...

    @abc.abstractmethod
    def close(self):
        """close resources"""
        ...


@dataclass
class FakeImage(ImageBase):
    """fake rasterio image"""

    driver: str
    width: int
    height: int
    count: int
    dtype: str
    crs: CRS
    transform: affine.Affine
    nodata: Union[float, int] = None

    def __post_init__(self):
        """post init hook"""
        fpath, _ = FakeGeoImage(
            xsize=self.width,
            ysize=self.height,
            num_bands=self.count,
            data_type=self.dtype,
            out_dir=Path("/vsimem"),
            crs=self.crs.to_epsg(),
            nodata=self.nodata or -1,
        ).create(transform=self.transform)
        self.handle = rasterio.open(fpath)

    def close(self):
        """close resources"""
        self.handle.close()


@dataclass
class FakeCog(ImageBase):
    """fake cloud optimized geotiff"""

    src_img: FakeImage
    dst_kwargs: Dict
    indexes: Optional[Sequence[int]] = None
    nodata: Optional[Union[int, float]] = None
    dtype: Optional[str] = None
    add_mask: bool = False
    overview_level: Optional[int] = None
    overview_resampling: str = "nearest"
    web_optimized: bool = False
    latitude_adjustment: bool = True
    resampling: str = "nearest"
    config: Optional[Dict] = None
    allow_intermediate_compression: bool = False
    forward_band_tags: bool = False
    quiet: bool = False
    temporary_compression: str = "DEFLATE"

    def __post_init__(self):
        """post init hook"""
        self.memfile = MemoryFile()
        cog_translate(
            source=self.src_img.handle,
            dst_path=self.memfile.name,
            dst_kwargs=self.dst_kwargs,
            indexes=self.indexes,
            nodata=self.nodata,
            dtype=self.dtype,
            add_mask=self.add_mask,
            overview_level=self.overview_level,
            overview_resampling=self.overview_resampling,
            web_optimized=self.web_optimized,
            latitude_adjustment=self.latitude_adjustment,
            resampling=self.resampling,
            in_memory=True,
            config=self.config,
            allow_intermediate_compression=self.allow_intermediate_compression,
            forward_band_tags=self.forward_band_tags,
            quiet=self.quiet,
            temporary_compression=self.temporary_compression,
        )
        self.handle = self.memfile.open()

    def close(self):
        """close resources"""
        self.handle.close()
        self.memfile.close()

cog_fixtures.cog module # type: ignore base image class support context management support context management post init hook close resources fake rasterio image post init hook close resources fake cloud optimized geotiff post init hook close resources