import os
import pandas as pd
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor
from src.logger import get_logger
from src.datasets.images.utils import read_image, verify


logger = get_logger(__name__)


def assert_conversations_format(df: pd.DataFrame) -> None:
    logger.info("Verifying conversations format...")
    required_columns = {"id", "image", "image_source", "image_shard", "conversations"}
    if set(df.columns) != required_columns:
        raise ValueError(f"DataFrame must contain (only!) columns: {required_columns}. Got: {set(df.columns)}")

    logger.info("Conversations format is valid!")


def assert_no_empty_conversations(df: pd.DataFrame) -> None:
    logger.info("Verifying that there are no empty conversations...")
    empty_conversations = df[df["conversations"].apply(lambda x: len(x) == 0)]
    if not empty_conversations.empty:
        raise ValueError(f"Found {len(empty_conversations)} empty conversations")

    logger.info("No empty conversations found!")


def assert_correct_conversations(df: pd.DataFrame) -> None:
    def check_conversation(conversation):
        image_tokens = 0
        for i, message in enumerate(conversation):
            # contains only 'from' and 'value' keys
            if set(message.keys()) != {"from", "value"}:
                raise ValueError(f"Message keys must be 'from' and 'value', got {message.keys()}")

            # 'from' is either 'human' or 'gpt' in correct order
            if i % 2 == 0 and message["from"] != "human":
                raise ValueError(f"Expected 'human' at position {i}, got {message['from']}")
            if i % 2 == 1 and message["from"] != "gpt":
                raise ValueError(f"Expected 'gpt' at position {i}, got {message['from']}")

            # 'value' is a non-empty string
            if not isinstance(message["value"], str) or not message["value"].strip():
                raise ValueError(f"'value' must be a non-empty string, got {message['value']}")

            if message["from"] == "human":
                image_tokens += message["value"].count("<image>")

        # exactly one image token in human messages
        if image_tokens != 1:
            raise ValueError(
                f"Expected exactly one <image> token in human messages, got {image_tokens}"
            )

    logger.info("Verifying conversations content...")
    for _, row in tqdm(df.iterrows(), total=len(df), leave=False):
        try:
            check_conversation(row["conversations"])
        except ValueError as e:
            raise ValueError(f"Error in conversation with id {row['id']}: {e}")
        
    logger.info("All conversations are correctly formatted!")


def assert_images_exist(df: pd.DataFrame, images_dir: str, workers: int = 32) -> None:
    logger.info("Verifying that all images can be opened...")
    def check(row):
        root_dir = os.path.join(images_dir, row["image_source"])
        image = read_image(root_dir, row["image"], row["image_shard"])
        if not verify(image):
            raise ValueError(f"Image {row['image']} in shard {row['image_shard']} cannot be opened")

    records = df.to_dict("records")
    with ThreadPoolExecutor(max_workers=workers) as executor:
        for _ in tqdm(executor.map(check, records), total=len(records), leave=False):
            pass

    logger.info("All images can be opened!")

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser("Verify conversations dataset")
    parser.add_argument("--input_path", type=str, required=True)
    parser.add_argument("--images_dir", type=str, required=True)
    parser.add_argument("--workers", type=int, default=32)

    args = parser.parse_args()

    logger.info(f"Loading conversations from {args.input_path}")
    conversations_df = pd.read_json(args.input_path, lines=args.input_path.endswith(".jsonl"))

    assert_conversations_format(conversations_df)
    assert_no_empty_conversations(conversations_df)
    assert_correct_conversations(conversations_df)
    assert_images_exist(conversations_df, args.images_dir, args.workers)
