import boto3
import os
import tempfile
import json
from tqdm import tqdm


def download_s3_folder(bucket_name, s3_folder, profile_name="default"):
    """
    Download all files from an S3 folder to a temporary directory

    Args:
        bucket_name (str): Name of the S3 bucket
        s3_folder (str): Path of the folder in S3 to download
        profile_name (str): AWS profile name stored in ~/.aws/credentials

    Returns:
        str: Path to the temporary directory where files were downloaded
    """
    # Create a temporary directory to store the downloaded files
    temp_dir = tempfile.mkdtemp()
    print(f"Created temporary directory: {temp_dir}")

    # Create a session using your AWS profile
    session = boto3.Session(profile_name=profile_name)
    s3_client = session.client("s3")

    # List all objects in the folder
    if not s3_folder.endswith("/"):
        s3_folder += "/"

    try:
        # Get list of all objects in the folder
        paginator = s3_client.get_paginator("list_objects_v2")
        pages = paginator.paginate(Bucket=bucket_name, Prefix=s3_folder)

        # First, collect all objects to download
        objects_to_download = []
        for page in pages:
            if "Contents" not in page:
                print(f"No objects found in {s3_folder}")
                continue

            for obj in page["Contents"]:
                # Get the relative path by removing the folder prefix
                s3_key = obj["Key"]

                # Skip if it's a "folder" (empty object that ends with '/')
                if s3_key.endswith("/"):
                    continue

                objects_to_download.append(s3_key)

        # Download each object with progress bar
        download_count = 0
        for s3_key in tqdm(objects_to_download, desc="Downloading files", unit="file"):
            # Create the local file path
            relative_path = s3_key[len(s3_folder) :]
            local_file_path = os.path.join(temp_dir, relative_path)

            # Create directory if it doesn't exist
            os.makedirs(os.path.dirname(local_file_path), exist_ok=True)

            # Download the file
            s3_client.download_file(bucket_name, s3_key, local_file_path)
            download_count += 1

        print(f"Successfully downloaded {download_count} files to {temp_dir}")
        return temp_dir

    except Exception as e:
        print(f"Error downloading from S3: {e}")
        return None


def merge_jsonl_files(directory, output_file=None):
    """
    Merge all JSONL files in the directory into a single JSONL file

    Args:
        directory (str): Directory containing JSONL files
        output_file (str, optional): Path for the merged output file.
                                    If None, creates a file in the same directory.

    Returns:
        str: Path to the merged file
    """
    if output_file is None:
        output_file = os.path.join(directory, "merged_output.jsonl")

    # Find all JSONL files in the directory and subdirectories
    jsonl_files = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.endswith(".jsonl"):
                jsonl_files.append(os.path.join(root, file))

    if not jsonl_files:
        print(f"No JSONL files found in {directory}")
        return None

    print(f"Found {len(jsonl_files)} JSONL files to merge")

    # Merge the files with progress bar
    line_count = 0
    with open(output_file, "w") as outfile:
        for jsonl_file in tqdm(jsonl_files, desc="Merging JSONL files", unit="file"):
            try:
                with open(jsonl_file, "r") as infile:
                    for line in infile:
                        # Validate that each line is valid JSON
                        try:
                            json.loads(line.strip())
                            outfile.write(line)
                            line_count += 1
                        except json.JSONDecodeError:
                            # Just log to stderr to not interfere with progress bar
                            import sys

                            print(
                                f"Skipping invalid JSON line in {jsonl_file}",
                                file=sys.stderr,
                            )
            except Exception as e:
                import sys

                print(f"Error processing {jsonl_file}: {e}", file=sys.stderr)

    print(f"Successfully merged {line_count} lines into {output_file}")
    return output_file


if __name__ == "__main__":
    # Example usage:
    bucket_name = "webdataset"
    s3_folder = "bundles/v0/discogs_subset_50k/paint_stems/metas"
    profile_name = "oracle"  # The profile name in ~/.aws/credentials
    output_file = (
        "merged_data.jsonl"  # Path where you want to save the final merged file
    )

    # Download the S3 folder
    download_dir = download_s3_folder(bucket_name, s3_folder, profile_name)

    if download_dir:
        print(f"Files are available at: {download_dir}")

        # Merge all JSONL files to the output location (outside temp dir)
        merged_file = merge_jsonl_files(download_dir, output_file)

        if merged_file:
            print(f"All JSONL files have been merged into: {merged_file}")

        # Clean up the temporary directory
        import shutil

        shutil.rmtree(download_dir)
        print(f"Temporary directory {download_dir} has been deleted")