Jesper Fjellin

# Begin transaction for all operations
try:
    conn.autocommit = False

    # Initialize rich Progress
    with Progress(
        SpinnerColumn(),
        TextColumn("[cyan]{task.description:<30}"),
        BarColumn(bar_width=30),
        "[progress.percentage]{task.percentage:>3.0f}%",
        TimeElapsedColumn(),
        console=console,
        expand=False
    ) as progress:
        task = progress.add_task("       Processing files", total=len(file_info_list))

        for info in file_info_list:
            file = info['file']
            table_name = args.table if args.table else info['table_name']
            qualified_table = f"{schema}.{table_name}"
            gdf = info['gdf']

            try:
                logger.info(f"Processing {file}")

                # Handle geometry column naming
                if gdf.geometry.name != 'geom':
                    gdf = gdf.rename_geometry('geom')
                    gdf.set_geometry('geom', inplace=True)
                    gdf.set_crs(gdf.crs, inplace=True)

                # Ensure valid CRS
                if not gdf.crs:
                    logger.warning(f"No CRS found in {file}, defaulting to EPSG:4326")
                    gdf.set_crs(epsg=4326, inplace=True)

                if args.table:
                    # When using --table, only keep the geometry column
                    gdf = gdf[['geom']]

                    if table_name not in existing_tables:
                        # Get SRID from data or args
                        srid = args.epsg if args.epsg else gdf.crs.to_epsg()
                        if not srid:
                            srid = 4326  # Default to WGS84 if no SRID found

                        # Create table with generic geometry type and SRID
                        if create_generic_geometry_table(conn, engine, table_name, srid, schema):
                            existing_tables.append(table_name)
                        else:
                            continue

                        # For new tables, all geometries are new
                        logger.info(f"Found {format(len(gdf), ',').replace(',', ' ')} [green]new[/] geometries.")

                        # Append first batch of geometries
                        logger.info(f"Appending {format(len(gdf), ',').replace(',', ' ')} geometries to '{qualified_table}'")
                        if append_geometries(conn, engine, gdf, table_name, schema):
                            total_new += len(gdf)
                    else:
                        # Compare geometries before appending
                        new_geoms, updated_geoms, identical_geoms = compare_geometries(
                            gdf, conn, table_name, 'geom', schema=schema, exclude_columns=[], args=args
                        )

                        # Create summary of differences for this dataset
                        num_new = len(new_geoms) if new_geoms is not None else 0
                        num_updated = len(updated_geoms) if updated_geoms is not None else 0
                        num_identical = len(identical_geoms) if identical_geoms is not None else 0

                        logger.info(f"Found {format(num_new, ',').replace(',', ' ')} [green]new[/] geometries, "
                                    f"{format(num_updated, ',').replace(',', ' ')} [yellow]updated[/] geometries, and "
                                    f"{format(num_identical, ',').replace(',', ' ')} [red]identical[/] geometries skipped.")

                        if new_geoms is not None and not new_geoms.empty:
                            logger.info(f"Appending {format(len(new_geoms), ',').replace(',', ' ')} geometries to '{qualified_table}'")
                            # Use schema parameter in to_postgis
                            new_geoms.to_postgis(
                                name=table_name,
                                con=engine,
                                schema=schema,
                                if_exists='append',
                                index=False
                            )
                            total_new += len(new_geoms)

                        if identical_geoms is not None:
                            total_identical += len(identical_geoms)
                elif table_name in existing_tables:
                    logger.info(f"Table {qualified_table} exists, analyzing differences...")

                    # Get common columns between GDF and database table
                    common_columns = [col for col in gdf.columns if col != 'geom']
                    if not common_columns:
                        # If no common columns, just compare geometries
                        columns_sql = "MD5(ST_AsBinary(geom)) as geom_hash"
                    else:
                        quoted_columns = ', '.join(f'"{col}"' for col in common_columns)
                        columns_sql = f"MD5(ST_AsBinary(geom)) as geom_hash, {quoted_columns}"

                    # Check existing geometry column name in PostGIS
                    existing_geom_col = get_db_geometry_column(conn, table_name, schema=schema)

                    # Only keep 'geometry' if it's the existing column name, otherwise use 'geom'
                    target_geom_col = 'geometry' if existing_geom_col == 'geometry' else 'geom'

                    if gdf.geometry.name != target_geom_col:
                        logger.debug(f"Renaming geometry column from '{gdf.geometry.name}' to '{target_geom_col}'")
                        gdf = gdf.rename_geometry(target_geom_col)
                        gdf.set_geometry(target_geom_col, inplace=True)
                        gdf.set_crs(gdf.crs, inplace=True)  # Preserve CRS

                    new_geoms, updated_geoms, identical_geoms = compare_geometries(
                        gdf, conn, table_name, gdf.geometry.name, schema=schema, exclude_columns=exclude_cols, args=args
                    )

                    # Create summary of differences
                    num_new = len(new_geoms) if new_geoms is not None else 0
                    num_updated = len(updated_geoms) if updated_geoms is not None else 0
                    num_identical = len(identical_geoms) if identical_geoms is not None else 0

                    total_new += num_new
                    total_updated += num_updated
                    total_identical += num_identical

                    logger.info(f"Found {format(num_new, ',').replace(',', ' ')} [green]new[/] geometries, "
                                f"{format(num_updated, ',').replace(',', ' ')} [yellow]updated[/] geometries, and "
                                f"{format(num_identical, ',').replace(',', ' ')} [red]identical[/] geometries skipped. "
                                "Skipping identical geometries...")

                    # Handle new geometries
                    if num_new > 0:
                        try:
                            # Use schema parameter in to_postgis
                            new_geoms.to_postgis(
                                name=table_name,
                                con=engine,
                                schema=schema,
                                if_exists='append',
                                index=False
                            )
                            logger.info(f"Successfully appended {num_new} new geometries to {qualified_table}")
                        except Exception as e:
                            logger.error(f"Error appending new geometries: {e}")
                    else:
                        pass

                    # Handle updated geometries (if implemented)
                    if num_updated > 0:
                        update_geometries(updated_geoms, table_name, engine, unique_id_column='osm_id')  # Adjust unique_id_column as needed

                else:
                    num_geometries = len(gdf)
                    logger.info(f"Found {num_geometries} new geometries to import into new table '{qualified_table}'")

                    # Add coordinate printing for new tables
                    if args.coordinates:
                        for idx, row in gdf.iterrows():
                            print_geometry_details(row, "NEW", args.coordinates)

                    try:
                        # Write to PostGIS with schema
                        gdf.to_postgis(
                            name=table_name,  # Use unqualified name
                            con=engine,
                            schema=schema,    # Specify schema separately
                            if_exists='replace',
                            index=False
                        )

                        # Verify the table was created
                        cursor = conn.cursor()
                        cursor.execute("""
                            SELECT EXISTS (
                                SELECT 1
                                FROM information_schema.tables
                                WHERE table_schema = %s
                                AND table_name = %s
                            );
                        """, (schema, table_name))
                        table_exists = cursor.fetchone()[0]
                        cursor.close()

                        if table_exists:
                            logger.info(f"Successfully imported {num_geometries} geometries to new table '{qualified_table}'")
                            create_spatial_index(conn, table_name, schema=schema, geom_column='geom')
                            existing_tables.append(table_name)
                            total_new += num_geometries
                        else:
                            logger.error(f"Failed to create table '{qualified_table}'")

                    except Exception as e:
                        logger.error(f"Error importing '{file}': {e}")
                        continue

            except Exception as e:
                logger.error(f"Error processing '{file}': {e}")
                continue

            progress.advance(task)

    # Commit everything at once
    conn.commit()
    logger.info("All changes committed successfully")
    logger.info("Summary of tasks:\n"
                f"{format(total_new, ',').replace(',', ' ')} [green]new[/] geometries added, "
                f"{format(total_updated, ',').replace(',', ' ')} [yellow]updated[/] geometries, "
                f"{format(total_identical, ',').replace(',', ' ')} [red]identical[/] geometries skipped")

except Exception as e:
    conn.rollback()  # Rollback the transaction on error
    logger.error(f"An error occurred: {e}. All changes have been rolled back.")
    # Optionally restore from backups here
finally:
    conn.autocommit = True  # Reset autocommit

def backup_tables(conn, tables, schema='public'):
    """Create file backups of all affected tables before processing."""
    timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
    backup_dir = os.path.join(os.getcwd(), 'backups')
    backup_info = {}
    
    for table in tables:
        if not check_table_exists(conn, table, schema):
            logger.info(f"Table '{schema}.{table}' does not exist, no backup needed.")
            continue
        
        backup_file = os.path.join(backup_dir, f"{table}_backup_{timestamp}.sql")
        
        try:
            # Create pg_dump command
            cmd = [
                'pg_dump',
                f'--host={conn.info.host}',
                f'--port={conn.info.port}',
                f'--username={conn.info.user}',
                f'--dbname={conn.info.dbname}',
                f'--table={schema}.{table}',
                '--format=p',
                f'--file={backup_file}'
            ]
            
            # Set PGPASSWORD environment variable for the subprocess
            env = os.environ.copy()
            env['PGPASSWORD'] = conn.info.password
            
            # Execute pg_dump
            subprocess.run(cmd, env=env, check=True, capture_output=True)
            
            backup_info[table] = backup_file
            logger.info(f"Created backup of '{schema}.{table}' to '{backup_file}'")
            
            # Manage old backups
            manage_old_backups(backup_dir, table)
            
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to backup table '{schema}.{table}': {e.stderr.decode()}")
            return None
            
    return backup_info

def process_files(args, conn, engine, existing_tables, schema):
    total_new = 0
    total_updated = 0
    total_identical = 0

    # Determine file extensions to process
    supported_extensions = ['.shp', '.geojson', '.json', '.gpkg', '.kml', '.gml']

    # Collect files to process (only in specified directory, not subdirectories)
    file_info_list = []

    # List files only in the specified directory
    for file in os.listdir(args.filepath):
        if any(file.lower().endswith(ext) for ext in supported_extensions):
            full_path = os.path.join(args.filepath, file)

            # Skip if not a file (e.g., if it's a directory)
            if not os.path.isfile(full_path):
                continue

            table_name = os.path.splitext(file)[0].lower()

            try:
                gdf = gpd.read_file(full_path)

                # Handle CRS
                source_crs = gdf.crs
                if args.epsg:
                    # User specified an EPSG
                    if source_crs and source_crs.to_epsg() != args.epsg:
                        logger.info(f"Reprojecting from EPSG:{source_crs.to_epsg()} to EPSG:{args.epsg}")
                        gdf.set_crs(source_crs, inplace=True)  # Ensure source CRS is set
                        gdf = gdf.to_crs(epsg=args.epsg)
                    else:
                        gdf.set_crs(epsg=args.epsg, inplace=True)
                elif not source_crs:
                    # No source CRS and no user EPSG specified, default to 4326
                    logger.warning(f"No CRS found in {file}, defaulting to EPSG:4326")
                    gdf.set_crs(epsg=4326, inplace=True)
                else:
                    # Keep source CRS
                    pass

                input_geom_col = gdf.geometry.name
                file_info_list.append({
                    'file': file,
                    'full_path': full_path,
                    'table_name': table_name,
                    'gdf': gdf,
                    'input_geom_col': input_geom_col
                })
            except Exception as e:
                logger.error(f"[red]Error reading '{file}': {e}[/red]")
                continue

    if not file_info_list:
        logger.warning("[red]No spatial files found to process.[/red]")
        return

def compare_geometries(gdf: GeoDataFrame, conn, table_name: str, geom_column: str = 'geom', schema: str = 'public', exclude_columns: List[str] = None, args=None):
cursor = conn.cursor()

# Get the actual geometry column name from the database
db_geom_column = get_db_geometry_column(conn, table_name, schema=schema)
if not db_geom_column:
    logger.error(f"No geometry column found in table '{schema}.{table_name}'")
    return None, None, None

sql = f"""
SELECT MD5(ST_AsBinary({db_geom_column})) as geom_hash
FROM {schema}.{table_name}
"""

# Get existing geometry hashes from database
existing_hashes = set()
with conn.cursor() as cur:
    cur.execute(sql)
    for row in cur.fetchall():
        existing_hashes.add(row[0])

# Create temporary copy of GDF for comparison
comparison_gdf = gdf.copy()
comparison_gdf['geom_hash'] = comparison_gdf[geom_column].apply(compute_geom_hash)

# Compare with database hashes
new_geometries = []
identical_geometries = []

for idx, row in comparison_gdf.iterrows():
    geom_hash = row['geom_hash']
    if geom_hash in existing_hashes:
        identical_geometries.append(row)
    else:
        new_geometries.append(row)

# Convert lists to GeoDataFrames
new_gdf = GeoDataFrame(new_geometries, geometry=geom_column, crs=gdf.crs) if new_geometries else GeoDataFrame(columns=gdf.columns)
identical_gdf = GeoDataFrame(identical_geometries, geometry=geom_column, crs=gdf.crs) if identical_geometries else GeoDataFrame(columns=gdf.columns)

# Remove temporary hash column
for gdf_temp in [new_gdf, identical_gdf]:
    if 'geom_hash' in gdf_temp.columns:
        gdf_temp.drop('geom_hash', axis=1, inplace=True)

return new_gdf if not new_gdf.empty else None, None, identical_gdf if not identical_gdf.empty else None
                

def update_geometries(gdf, table_name, engine, unique_id_column):
"""Update existing geometries in PostGIS table."""
if gdf is None or gdf.empty:
return

try:
# Create temporary table for updates
temp_table = f"temp_{table_name}"
gdf.to_postgis(temp_table, engine, if_exists='replace', index=False)

with engine.connect() as connection:
    from sqlalchemy import text
    
    # First, check for and add any new columns
    cursor = connection.execute(text(f"""
        SELECT column_name 
        FROM information_schema.columns 
        WHERE table_schema = 'public' AND table_name = '{table_name}'
    """))
    existing_columns = {row[0] for row in cursor}
    
    # Get new columns from the GeoDataFrame
    new_columns = set(gdf.columns) - existing_columns
    
    # Add any new columns to the main table
    for col in new_columns:
        # Determine column type from GeoDataFrame
        dtype = gdf[col].dtype
        if dtype == 'object':
            sql_type = 'TEXT'
        elif dtype == 'int64':
            sql_type = 'INTEGER'
        elif dtype == 'float64':
            sql_type = 'DOUBLE PRECISION'
        else:
            sql_type = 'TEXT'  # Default to TEXT for unknown types
        
        logger.info(f"Adding new column '{col}' with type {sql_type}")
        connection.execute(text(f'ALTER TABLE "{table_name}" ADD COLUMN IF NOT EXISTS "{col}" {sql_type}'))
    
    # Now proceed with the update
    columns = [col for col in gdf.columns if col != unique_id_column]
    update_cols = ", ".join([f'"{col}" = s."{col}"' for col in columns])
    
    sql = text(f"""
        UPDATE "{table_name}" t
        SET {update_cols}
        FROM "{temp_table}" s
        WHERE t.{unique_id_column} = s.{unique_id_column}
    """)
    logger.debug(f"Executing update SQL: {sql}")
    connection.execute(sql)
    connection.execute(text(f'DROP TABLE IF EXISTS "{temp_table}"'))
    connection.commit()

logger.info(f"Successfully updated {len(gdf)} geometries in {table_name}")
except Exception as e:
logger.error(f"Error updating geometries: {e}")
                

# Collect geometry column names
geom_col_files = defaultdict(list)
for info in file_info_list:
    geom_col_files[info['input_geom_col']].append(info['file'])

# Handle geometry column renaming
for info in file_info_list:
    gdf = info['gdf']
    if gdf.geometry.name != 'geom':
        gdf = gdf.rename_geometry('geom')
        gdf.set_geometry('geom', inplace=True)
        gdf.set_crs(gdf.crs, inplace=True)  # Preserve CRS
        info['gdf'] = gdf
        info['input_geom_col'] = 'geom'
                

def check_crs_compatibility(gdf, conn, table_name, geom_column, args):
cursor = conn.cursor()
# Check if the table exists
cursor.execute("""
    SELECT EXISTS (
        SELECT 1
        FROM information_schema.tables 
        WHERE table_schema = 'public' AND table_name = %s
    );
""", (table_name,))
table_exists = cursor.fetchone()[0]

if not table_exists:
    # Table does not exist, proceed without CRS check
    cursor.close()
    return gdf  # Proceed with the current GeoDataFrame

# Table exists, retrieve existing SRID using ST_SRID
try:
    cursor.execute(f"""
        SELECT ST_SRID("{geom_column}") FROM "{table_name}" WHERE "{geom_column}" IS NOT NULL LIMIT 1;
    """)
    result = cursor.fetchone()
    if result:
        existing_srid = result[0]
        logger.info(f"Existing SRID for '{table_name}' is {existing_srid}")
    else:
        logger.warning(f"No geometries found in '{table_name}' to determine SRID")
        existing_srid = None
except Exception as e:
    logger.error(f"Error retrieving SRID for '{table_name}': {e}")
    conn.rollback()
    cursor.close()
    return None  # Skip this file due to error

# Get the SRID of the new data
new_srid = gdf.crs.to_epsg()
if new_srid is None:
    logger.warning(f"No EPSG code found for the CRS of the new data for '{table_name}'")
    if args.overwrite:
        action = 'y'
    else:
        action = console.input(f"Proceed without CRS check for '{table_name}'? (y/n): ")
    if action.lower() != 'y':
        logger.info(f"Skipping '{table_name}' due to unknown CRS")
        cursor.close()
        return None  # Skip this file
else:
    logger.info(f"CRS of new data for '{table_name}' is EPSG:{new_srid}")

# Compare SRIDs
if existing_srid and new_srid != existing_srid:
    logger.warning(f"CRS mismatch for '{table_name}': Existing SRID {existing_srid}, New SRID {new_srid}")
    if args.overwrite:
        action = 'y'
    else:
        action = console.input(f"Reproject new data to SRID {existing_srid}? (y/n): ")
    if action.lower() == 'y':
        try:
            gdf = gdf.to_crs(epsg=existing_srid)
            logger.info(f"Reprojected new data to SRID {existing_srid}")
        except Exception as e:
            logger.error(f"[red]Error reprojecting data for '{table_name}': {e}[/red]")
            cursor.close()
            return None  # Skip this file due to reprojection error
    else:
        logger.info(f"Skipping '{table_name}' due to CRS mismatch")
        cursor.close()
        return None  # Skip this file
else:
    logger.info(f"CRS is compatible for '{table_name}'")

cursor.close()
return gdf  # Return the (possibly reprojected) GeoDataFrame
                

def create_spatial_index(conn, table_name, schema='public', geom_column='geom'):
cursor = conn.cursor()
# Get the actual geometry column name
actual_geom_column = get_db_geometry_column(conn, table_name, schema=schema) or geom_column
index_name = f"{schema}_{table_name}_{actual_geom_column}_idx"

try:
    cursor.execute(f"""
        CREATE INDEX IF NOT EXISTS "{index_name}"
        ON "{schema}"."{table_name}"
        USING GIST ("{actual_geom_column}");
    """)
    conn.commit()
    logger.info(f"Spatial index created on table '{schema}.{table_name}'")
except Exception as e:
    logger.error(f"Error creating spatial index on '{schema}.{table_name}': {e}")
    conn.rollback()
finally:
    cursor.close()

// Parallel processing implementation
use rayon::prelude::*;  // Parallel processing functionality
use std::collections::HashMap;

let cell_size = buffer_distance * 2.0;
let buffer_distance_squared = buffer_distance * buffer_distance;

// Create spatial grid for efficient neighbor searching
let mut grid: HashMap<(i32, i32), Vec<(usize, (f64, f64))>> = 
    HashMap::with_capacity(points.len() / 10);

// Assign points to grid cells
for (i, &(x, y)) in points.iter().enumerate() {
    let cell = (
        (x / cell_size).floor() as i32,
        (y / cell_size).floor() as i32,
    );
    grid.entry(cell)
        .or_insert_with(Vec::new)
        .push((i, (x, y)));
}

// Process grid cells in parallel
let results: Vec<(usize, usize)> = grid.par_iter()
    .flat_map(|(&(cell_x, cell_y), cell_points)| {
        let mut local_results = Vec::new();
        
        // Check current and neighboring cells
        let neighbor_cells = [
            (cell_x, cell_y),
            (cell_x, cell_y + 1),
            (cell_x + 1, cell_y - 1),
            (cell_x + 1, cell_y),
            (cell_x + 1, cell_y + 1),
        ];
        
        for &neighbor_cell in neighbor_cells.iter() {
            if let Some(neighbor_points) = grid.get(&neighbor_cell) {
                for &(i, (x1, y1)) in cell_points {
                    for &(j, (x2, y2)) in neighbor_points {
                        if i < j {
                            let dx = x1 - x2;
                            let dy = y1 - y2;
                            let distance_squared = dx * dx + dy * dy;
                            if distance_squared <= buffer_distance_squared && distance_squared > 0.0 {
                                local_results.push((i, j));
                            }
                        }
                    }
                }
            }
        }
        local_results
    })
    .collect();
                    

// Python binding setup
use pyo3::prelude::*;
use std::collections::HashMap;
use rayon::prelude::*;

#[pyfunction]
fn points_within_buffer_batch(
    points: Vec<(f64, f64)>, 
    buffer_distance: f64
) -> PyResult> {
    let cell_size = buffer_distance * 2.0;
    let buffer_distance_squared = buffer_distance * buffer_distance;
    
    let mut grid: HashMap<(i32, i32), Vec<(usize, (f64, f64))>> = 
        HashMap::with_capacity(points.len() / 10);
    
    // Build spatial grid
    for (i, &(x, y)) in points.iter().enumerate() {
        let cell = (
            (x / cell_size).floor() as i32,
            (y / cell_size).floor() as i32,
        );
        grid.entry(cell)
            .or_insert_with(Vec::new)
            .push((i, (x, y)));
    }
    
    // Process points in parallel and find pairs within buffer distance
    let results = grid.par_iter()
        .flat_map(|(&(cell_x, cell_y), cell_points)| {
            let mut local_results = Vec::new();
            let neighbor_cells = [
                (cell_x, cell_y),
                (cell_x, cell_y + 1),
                (cell_x + 1, cell_y - 1),
                (cell_x + 1, cell_y),
                (cell_x + 1, cell_y + 1),
            ];
            
            for &neighbor_cell in neighbor_cells.iter() {
                if let Some(neighbor_points) = grid.get(&neighbor_cell) {
                    for &(i, (x1, y1)) in cell_points {
                        for &(j, (x2, y2)) in neighbor_points {
                            if i < j {
                                let dx = x1 - x2;
                                let dy = y1 - y2;
                                let distance_squared = dx * dx + dy * dy;
                                if distance_squared <= buffer_distance_squared && distance_squared > 0.0 {
                                    local_results.push((i, j));
                                }
                            }
                        }
                    }
                }
            }
            local_results
        })
        .collect();

    Ok(results)
}

#[pymodule]
fn rust_bindings(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(points_within_buffer_batch, m)?)?;
    Ok(())
}