import hashlib
import os
import json
from datasets import load_dataset
from tqdm import tqdm

OUTPUT_DIR = "data/flickr8k"

def download_flickr8k(output_dir=OUTPUT_DIR):
    """
    Downloads the Flickr8k dataset from HuggingFace and saves it in a format
    compatible with the toan toolkit (images folder + JSON registry).
    """
    images_dir = os.path.join(output_dir, "images")
    os.makedirs(images_dir, exist_ok=True)

    try:
        ds = load_dataset("Naveengo/flickr8k", trust_remote_code=False)
    except Exception as e:
        print(f"Error loading 'Naveengo/flickr8k': {e}")
        return

    registry = {}
    
    for split in ds.keys():
        for item in tqdm(ds[split]):
            img_obj = item['image']
            
            if 'filename' in item and item['filename']:
                fname = item['filename']
                if not fname.lower().endswith('.jpg'):
                    fname += ".jpg"
            elif 'image_id' in item and item['image_id']:
                fname = f"{item['image_id']}.jpg"
            else:
                # Fallback: Hash the image content to generate unique ID
                # This ensures we aggregate captions for the same image
                # Important: Ensure image is in RGB mode for consistent bytes, though usually it is.
                if img_obj.mode != 'RGB':
                    img_obj = img_obj.convert('RGB')
                
                img_bytes = img_obj.tobytes()
                img_hash = hashlib.md5(img_bytes).hexdigest()
                fname = f"{img_hash}.jpg"

            save_path = os.path.join(images_dir, fname)

            if fname not in registry:
                if not os.path.exists(save_path):
                    img_obj.save(save_path)
                
                registry[fname] = {
                    "image_path": os.path.join("images", fname),
                    "captions": [],
                    "split": split
                }
            

            caption = item.get('caption') or item.get('text') or []
            if isinstance(caption, str):
                registry[fname]["captions"].append(caption)
            elif isinstance(caption, list):
                registry[fname]["captions"].extend(caption)


    json_path = os.path.join(output_dir, "dataset.json")
    with open(json_path, 'w') as f:
        json.dump(registry, f, indent=2)
    
    print(f"Flickr8k download complete. Registry saved to {json_path}")

if __name__ == "__main__":
    download_flickr8k()