import unittest
from unittest.mock import patch, MagicMock
import os
from geodata_pipeline.heightmaps import export_heightmaps
from geodata_pipeline.config import Config

class TestHeightmapsOptimization(unittest.TestCase):
    @patch("geodata_pipeline.heightmaps.gdal.Warp")
    @patch("geodata_pipeline.heightmaps.gdal.Translate")
    @patch("geodata_pipeline.heightmaps.open_dataset")
    @patch("geodata_pipeline.heightmaps.build_vrt")
    @patch("geodata_pipeline.heightmaps.ensure_dir")
    @patch("geodata_pipeline.heightmaps.ensure_parent")
    @patch("geodata_pipeline.heightmaps._compute_ds_minmax")
    @patch("geodata_pipeline.heightmaps.cleanup_aux_files")
    @patch("geodata_pipeline.heightmaps.glob.glob")
    @patch("builtins.open")
    def test_export_heightmaps_in_memory_warp(self, mock_open_file, mock_glob, mock_cleanup, mock_minmax, mock_ensure_p, mock_ensure_d, mock_build_vrt, mock_open_ds, mock_translate, mock_warp):
        # Setup mocks
        mock_glob.return_value = ["tile1.tif"]
        mock_open_ds.return_value.GetRasterBand.return_value.ComputeRasterMinMax.return_value = (0, 100)
        mock_minmax.return_value = (10, 90, 100) # min, max, count
        
        # Mock VRT dataset
        mock_ds = MagicMock()
        mock_open_ds.return_value = mock_ds
        mock_ds.GetRasterBand.return_value.ComputeRasterMinMax.return_value = (0, 100)
        
        # Mock individual tile dataset for GT
        mock_tds = MagicMock()
        mock_tds.GetGeoTransform.return_value = (1000, 1, 0, 2000, 0, -1)
        mock_tds.RasterXSize = 1000
        mock_tds.RasterYSize = 1000
        
        # Side effect for open_dataset: first VRT, then tile1.tif
        mock_open_ds.side_effect = [mock_ds, mock_tds]
        
        cfg = Config.default()
        cfg.raw.dgm1_dir = "raw/dgm1"
        cfg.work.work_dir = "work"
        cfg.export.heightmap_dir = "export"
        cfg.export.manifest_path = "manifest.csv"
        
        # Mock Warp to return a memory dataset
        mock_mem_ds = MagicMock()
        mock_warp.return_value = mock_mem_ds
        
        export_heightmaps(cfg)
        
        # Verify Warp was called with empty string or MEM driver destination
        warp_args, warp_kwargs = mock_warp.call_args
        # Warp(destNameOrDestDS, srcDSOrSrcDSTab, ...)
        dest = warp_args[0]
        # In the original code, dest was a path. We want to check it's changed to "" for MEM.
        # But wait, my refactor will use "" for destName.
        
        # Verify Translate used the warped dataset
        translate_args, translate_kwargs = mock_translate.call_args
        # Translate(destName, srcDS, ...)
        src_ds = translate_args[1]
        self.assertEqual(src_ds, mock_mem_ds)

if __name__ == "__main__":
    unittest.main()