from datasets import load_dataset, Image
import os
import pandas as pd
from tqdm import tqdm
from src.logger import get_logger
from src.config import ProjectConfig

logger = get_logger(__name__)

project_config = ProjectConfig()

def get_xm3600_data(language: str) -> None:
    metadata_path = os.path.join(project_config.data_dir, f"raw/xm3600_{language}")
    images_path = os.path.join(project_config.data_dir, f"static/raw/xm3600_{language}")

    logger.info(f"Metadata path: {metadata_path}")
    logger.info(f"Images path: {images_path}")

    dataset = load_dataset("floschne/xm3600", split=language)
    dataset = dataset.cast_column("image", Image(decode=True))

    os.makedirs(images_path, exist_ok=True)

    rows = []

    for row in tqdm(dataset):
        img_id = row["image_id"]
        img = row["image"]
        caption = row["captions"]
        final_images_path = os.path.join(images_path, f"{img_id}.jpg")
        img.save(final_images_path)
        rows.append({"image_id": img_id, "image_path": final_images_path, "caption_pl": caption})

    df = pd.DataFrame(rows)
    csv_path = f"{metadata_path}/captions_{language}.csv"
    df.to_csv(csv_path, index=False, encoding="utf-8")
    
    logger.info(f"Saved {len(df)} entries.")
    logger.info(f"Images: {images_path}")
    logger.info(f"Metadata: {csv_path}")
    
if __name__ == "__main__":
    
    logger.info("Getting the data for xm3600 dataset..")
    language = "pl"
    get_xm3600_data(language)
