import os
import shutil
from osgeo import gdal

def fix_manual_masks(raw_dir, work_dir):
    manual_files = [f for f in os.listdir(raw_dir) if f.endswith(".png")]
    
    for manual_file in manual_files:
        manual_path = os.path.join(raw_dir, manual_file)
        
        parts = manual_file.replace("_viz.png", "").split("_")
        if len(parts) >= 4:
            tile_prefix = f"{parts[0]}_{parts[1]}_{parts[2]}_{parts[3]}"
            work_file = f"{tile_prefix}_1_rp_mask.png"
            work_path = os.path.join(work_dir, work_file)
            
            if not os.path.exists(work_path):
                print(f"Skipping {manual_file}: Work file not found.")
                continue
                
            print(f"Fixing {manual_file} using {work_file}...")
            
            ds_work = gdal.Open(work_path)
            gt_work = ds_work.GetGeoTransform()
            proj_work = ds_work.GetProjection()
            target_w = ds_work.RasterXSize
            target_h = ds_work.RasterYSize
            
            minx = gt_work[0]
            maxy = gt_work[3]
            maxx = minx + (gt_work[1] * target_w)
            miny = maxy + (gt_work[5] * target_h)
            
            # Open Source as ReadOnly, then CreateCopy to MEM to edit
            ds_raw = gdal.Open(manual_path)
            mem_driver = gdal.GetDriverByName("MEM")
            ds_manual = mem_driver.CreateCopy("", ds_raw)
            ds_raw = None # Close file
            
            src_w = ds_manual.RasterXSize
            src_h = ds_manual.RasterYSize
            
            src_res_x = (maxx - minx) / src_w
            src_res_y = (miny - maxy) / src_h
            
            src_gt = (minx, src_res_x, 0, maxy, 0, src_res_y)
            ds_manual.SetGeoTransform(src_gt)
            ds_manual.SetProjection(proj_work)
            
            warp_options = gdal.WarpOptions(
                format="PNG",
                width=target_w,
                height=target_h,
                outputBounds=(minx, miny, maxx, maxy),
                resampleAlg=gdal.GRA_NearestNeighbour
            )
            
            temp_path = manual_path + ".tmp.png"
            ds_fixed = gdal.Warp(temp_path, ds_manual, options=warp_options)
            ds_fixed = None 
            ds_manual = None
            ds_work = None
            
            shutil.move(temp_path, manual_path)
            
            work_wld = work_path.replace(".png", ".wld")
            if os.path.exists(work_wld):
                shutil.copy(work_wld, manual_path.replace(".png", ".wld"))
                
            print(f"  -> Fixed.")

if __name__ == "__main__":
    fix_manual_masks("raw/water_masks", "work/river_masks")