import os
import pandas as pd
import json
import fcntl


def get_and_save_individual_results(metric_name, get_results_func, output_dir):
    """
    Generalized function to save evaluation results to a CSV file.
    """

    # replace ending "/" of the path
    if output_dir[-1] == "/":
        output_dir = output_dir[:-1]

    csv_path = output_dir + f'_{metric_name}_results.csv'
    os.makedirs(os.path.dirname(csv_path), exist_ok=True)

    # Check if the results file already exists
    if not os.path.exists(csv_path):
        os.makedirs(os.path.dirname(csv_path), exist_ok=True)

        # Generate the results using the provided get_results_func
        results = get_results_func()

        # Ensure the directory exists
        if not os.path.isdir(os.path.dirname(csv_path)):
            os.makedirs(os.path.dirname(csv_path), exist_ok=True)

        # Convert the results to DataFrame
        results_df = pd.DataFrame(results)

        # Save the DataFrame to a CSV file
        results_df.to_csv(csv_path, index=False)
        print(f"Saved {metric_name} evaluation results for model to {csv_path}")

        # Print a sample of the results
        print("\nSample Results:")
        print(results_df.head())
    else:
        # If the file already exists, simply load the CSV into DataFrame
        results_df = pd.read_csv(csv_path)
        print(f"Loaded existing {metric_name} evaluation results for model from {csv_path}")

    return results_df


def update_summary_with_lock(summary_df, full_model_name, summary_dir, metric_name):
    """
    Updates the summary JSON file with the provided summary dataframe.
    Uses file locking to avoid race conditions when multiple processes access the file.
    """

    # Define the path for the summary JSON file
    summary_json_path = os.path.join(summary_dir, f'{metric_name}_{full_model_name.split("/")[0]}_summary.json')

    # Ensure the directory exists
    os.makedirs(os.path.dirname(summary_json_path), exist_ok=True)

    # Open the file with write access and lock it using fcntl
    with open(summary_json_path, 'a+') as f:
        # Lock the file
        fcntl.flock(f, fcntl.LOCK_EX)  # Exclusive lock

        # Check if the file is empty (i.e., does not exist or is newly created)
        f.seek(0)  # Move to the start of the file to check for its contents
        if f.read(1):  # If the file is not empty
            f.seek(0)  # Go back to the start of the file
            summary_json = json.load(f)
        else:
            summary_json = {}  # Initialize an empty dictionary if the file doesn't exist

        # Convert the summary dataframe to a dictionary
        summary_dict = summary_df.to_dict(orient='records')

        # Example: Update the specific metric for the model
        summary_json[full_model_name] = summary_dict

        # Move the file pointer to the start of the file before overwriting it
        f.seek(0)
        f.truncate()  # Clear the file content
        json.dump(summary_json, f, indent=4)

        # Unlock the file when done
        fcntl.flock(f, fcntl.LOCK_UN)  # Unlock the file


def determine_grid_dimension(total_size, image_size, padding_threshold):
    """
    Determines the number of images (columns/rows) and padding between them.

    Parameters:
    - total_size (int): Total width or height of the grid.
    - image_size (int): Size of each individual image.
    - padding_threshold (int): Maximum allowed padding between images.

    Returns:
    - (int, int): Number of images and padding size.
    """
    # Start from the maximum possible number of images
    max_n = (total_size + padding_threshold) // (image_size + padding_threshold)
    for n in range(max_n, 0, -1):
        if n == 1:
            padding = 0
            if total_size == image_size:
                return n, padding
            else:
                continue
        padding = (total_size - n * image_size) / (n - 1)
        if 0 <= padding <= padding_threshold:
            return n, int(round(padding))
    raise ValueError(f"Unable to determine grid dimension for total_size={total_size}")


def split_grid(image, image_size=512, padding_threshold=10, verbose=False):
    """
    Splits a grid image into individual images of size `image_size` x `image_size`.

    Parameters:
    - image (PIL.Image): The grid image to split.
    - image_size (int): The size of each individual image (default is 512).
    - padding_threshold (int): Maximum allowed padding between images (default is 10 pixels).

    Returns:
    - List[PIL.Image]: A list of individual images extracted from the grid.
    """
    grid_width, grid_height = image.size

    # Determine number of columns and padding_x
    n_cols, padding_x = determine_grid_dimension(grid_width, image_size, padding_threshold)

    # Determine number of rows and padding_y
    n_rows, padding_y = determine_grid_dimension(grid_height, image_size, padding_threshold)

    # Debugging Information
    if verbose:
        print(f"Grid Layout: {n_rows} rows x {n_cols} columns")
        print(f"Padding between images: x={padding_x} pixels, y={padding_y} pixels")

    individual_images = []

    for row in range(n_rows):
        for col in range(n_cols):
            left = col * (image_size + padding_x)
            upper = row * (image_size + padding_y)
            right = left + image_size
            lower = upper + image_size

            # Ensure the crop box is within the image boundaries
            right = min(right, grid_width)
            lower = min(lower, grid_height)

            # Crop the individual image
            individual_image = image.crop((left, upper, right, lower))
            individual_images.append(individual_image)

    return individual_images


def parse_from_images_dir(images_dir):
    path_parts = images_dir.split(os.sep)
    images_index = path_parts.index("images")
    prompt_file_name = path_parts[images_index + 1]
    full_model_name = os.path.join(*path_parts[images_index + 2:])

    results_dir = images_dir.replace("images", "metrics")
    summary_dir = os.path.join(*path_parts[:images_index + 2]).replace("images", "metrics")

    return prompt_file_name, full_model_name, results_dir, summary_dir
