import pandas as pd
import numpy as np


def upsert_csv(
    file_path, new_data, num_key_columns=5, target_columns=None, reset=False
):
    print("FILEPATH: ", file_path)
    # Convert new_data to a list if it's a single dictionary
    if isinstance(new_data, dict):
        new_data = [new_data]

    # Read existing CSV file or create a new DataFrame
    try:
        # Read CSV with strings for all columns to prevent auto-conversion
        df = pd.read_csv(file_path, dtype=str, index_col=False)
        # Convert 'True' and 'False' strings to boolean
        for col in df.columns:
            if set(df[col].unique()) == {"True", "False"}:
                df[col] = df[col].map({"True": True, "False": False})
    except FileNotFoundError:
        df = pd.DataFrame(columns=list(new_data[0].keys()))

    if reset:
        df = pd.DataFrame(columns=list(new_data[0].keys()))

    # Convert new_data to DataFrame
    new_df = pd.DataFrame(new_data)

    # Define key columns
    key_columns = df.columns[:num_key_columns].tolist()

    if target_columns:
        # Ensure target_columns is a list
        if isinstance(target_columns, str):
            target_columns = [target_columns]

        # Add new columns to df if they don't exist
        for col in target_columns:
            if col not in df.columns:
                df[col] = np.nan

        # Multiple column update logic
        for item in new_data:
            mask = np.all(
                df[key_columns].astype(str)
                == pd.DataFrame([item])[key_columns].astype(str).values,
                axis=1,
            )

            if mask.any():
                for col in target_columns:
                    if col in item:
                        df.loc[mask, col] = item[col]
            else:
                new_row = pd.DataFrame([item]).reindex(
                    columns=df.columns, fill_value=np.nan
                )
                df = pd.concat([df, new_row], ignore_index=True)
    else:
        # Multiple row update logic (unchanged)
        combined_df = pd.concat([df, new_df], ignore_index=True)
        combined_df = combined_df.drop_duplicates(
            subset=key_columns, keep="last"
        )
        df = combined_df

    # Convert boolean columns to 'True' and 'False' strings before saving
    for col in df.columns:
        if df[col].dtype == bool:
            df[col] = df[col].map({True: "True", False: "False"})

    df.to_csv(file_path, index=False)


if __name__ == "__main__":
    filename = "test.csv"
    data = [
        {
            "col1": "value1",
            "col2": "value2",
            "col3": "value3",
            "col4": "value4",
        }
    ]
    upsert_csv(filename, data)

    data = [
        {
            "col1": "1",
            "col2": "2",
            "col3": "3",
            "col4": "new2",
        },
    ]
    upsert_csv(filename, data, target_column="col4")
