import os
import tarfile
from tqdm import tqdm
import pandas as pd
from src.logger import get_logger
from PIL import Image

logger = get_logger(__name__)

ALLOWED_EXTENSIONS = (".png", ".jpg", ".jpeg", ".webp")


def list_images(images_dir: str) -> list[str]:
    images = []
    for root, _, files in os.walk(images_dir):
        for filename in files:
            if not filename.lower().endswith(ALLOWED_EXTENSIONS):
                continue

            full_path = os.path.join(root, filename)
            relative_path = os.path.relpath(full_path, images_dir)
            images.append(relative_path)

    return images


def split_into_shards(items: list, num_shards: int) -> dict[str, list[str]]:
    shards = {f"{i:05d}.tar": [] for i in range(num_shards)}
    for item in items:
        shard_id = hash(item) % num_shards
        shard_name = f"{shard_id:05d}.tar"
        shards[shard_name].append(item)
    return shards


def write_shards(shards: dict[str, list[str]], input_images_dir: str, output_images_dir: str):
    import concurrent.futures

    os.makedirs(output_images_dir, exist_ok=True)

    def verify_image(image_path: str) -> bool:
        try:
            with Image.open(image_path) as img:
                img.verify()  # Verify that it's a valid image
            return True
        except Exception as e:
            logger.warning(f"Invalid image {image_path}: {e}")
            return False

    def write_shard(shard_id: str, images: list[str]):
        shard_path = os.path.join(output_images_dir, shard_id)
        with tarfile.open(shard_path, "w") as tar:
            for image in tqdm(
                images, desc=f"Writing images to {shard_id}", leave=False, disable=True
            ):
                image_path = os.path.join(input_images_dir, image)
                assert os.path.exists(image_path), f"Image not found: {image_path}"
                if verify_image(image_path):
                    tar.add(image_path, arcname=image)

    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
        futures = [
            executor.submit(write_shard, shard_id, images) for shard_id, images in shards.items()
        ]
        for f in tqdm(
            concurrent.futures.as_completed(futures), total=len(futures), desc="Writing shards"
        ):
            f.result()


def create_shards(
    input_images_dir: str, output_images_dir: str, images_per_shard: int = 512
) -> pd.DataFrame:
    logger.info(f"Listing images in {input_images_dir}")
    images = list_images(input_images_dir)
    logger.info(f"Total images: {len(images)}")

    num_shards = max(1, len(images) // images_per_shard)
    shards = split_into_shards(images, num_shards)
    logger.info(f"Total shards: {len(shards)}")

    write_shards(shards, input_images_dir, output_images_dir)
    logger.info(f"Sharded images written to {output_images_dir}")

    return pd.DataFrame(
        [(shard_id, img) for shard_id, imgs in shards.items() for img in imgs],
        columns=["shard", "image"],
    )


def assign_shards(
    conversations_df: pd.DataFrame, shard_index: pd.DataFrame, remove_without_image: bool
) -> pd.DataFrame:
    image_to_shard = dict(zip(shard_index["image"], shard_index["shard"]))

    conversations_df["image_shard"] = conversations_df.apply(
        lambda row: image_to_shard.get(
            row["image"], row["image_shard"] if "image_shard" in row else None
        ),
        axis=1,
    )

    any_without_image = conversations_df["image_shard"].isnull().sum() > 0
    if any_without_image:
        missing_images = conversations_df[conversations_df["image_shard"].isnull()][
            "image"
        ].unique()
        logger.warning(f"Some images are missing in the shard index: {len(missing_images)}")
        logger.warning(f"Examples of missing images: {missing_images[:5]}")

    if any_without_image and remove_without_image:
        conversations_df = conversations_df[
            ~conversations_df["image_shard"].isnull()
        ].reset_index(drop=True)
        logger.info(f"Removed conversations without images. Remaining: {len(conversations_df)}")

    return conversations_df


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--input_images_dir", type=str, required=True)
    parser.add_argument("--create", action="store_true", help="Whether to create shards.")
    parser.add_argument("--output_images_dir", type=str, default=None)
    parser.add_argument("--images_per_shard", type=int, default=512)
    parser.add_argument(
        "--assign", action="store_true", help="Whether to assign shards to conversations."
    )
    parser.add_argument("--input_path", type=str, default=None)
    parser.add_argument("--output_path", type=str, default=None)
    parser.add_argument(
        "--remove_without_image",
        action="store_true",
        help="Whether to remove conversations without images.",
    )
    args = parser.parse_args()

    if args.create:
        logger.info("Creating image shards.")
        index_df = create_shards(
            args.input_images_dir, args.output_images_dir, args.images_per_shard
        )

        index_path = os.path.join(args.output_images_dir, "index.jsonl")
        index_df.to_json(index_path, orient="records", lines=True)
        logger.info(f"Index file written to {index_path}")

    if args.assign:
        logger.info("Assigning shards to conversations.")
        index_path = os.path.join(args.input_images_dir, "index.jsonl")
        assert os.path.exists(index_path), f"Index file not found: {index_path}"

        index_df = pd.read_json(index_path, lines=True)
        conversations_df = pd.read_json(args.input_path, lines=args.input_path.endswith(".jsonl"))

        # assign shards
        conversations_df = assign_shards(conversations_df, index_df, args.remove_without_image)

        conversations_df.to_json(
            args.output_path, orient="records", lines=args.output_path.endswith(".jsonl")
        )
        logger.info(f"Conversations with assigned shards written to {args.output_path}")
