import random
import pandas as pd
import json
from src.datasets.utils import add_image_token
from src.datasets.wit.utils import get_image_name
from src.logger import get_logger


logger = get_logger(__name__)


def apply_conversations_format(
    df: pd.DataFrame, instructions: list[str], seed: int = 42
) -> pd.DataFrame:
    random.seed(seed)
    logger.info("Applying conversations format")

    def get_id(row: pd.Series) -> str:
        # prevent negative values
        return "wit_" + str(hash((row["page_url"], row["image_url"])) & 0xFFFFFFFFFFFFFFFF)

    def get_conversation(caption: str, image_token: str = "<image>") -> list[dict[str, str]]:
        return [
            {
                "from": "human",
                "value": add_image_token(
                    random.choice(instructions),
                    random.choice(["preceding", "following"]),
                    image_token=image_token,
                ),
            },
            {"from": "gpt", "value": caption},
        ]

    conversations_df = pd.DataFrame()
    conversations_df["id"] = df.apply(get_id, axis=1)
    conversations_df["image"] = df["image_url"].apply(get_image_name)
    conversations_df["image_source"] = "wit"
    conversations_df["conversations"] = df["caption_reference_description"].apply(
        get_conversation
    )

    return conversations_df


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--input_path", type=str, required=True)
    parser.add_argument("--instructions_path", type=str, required=True)
    parser.add_argument("--output_path", type=str, required=True)
    parser.add_argument("--seed", type=int, default=42)

    args = parser.parse_args()

    logger.info(f"Reading wit data from {args.input_path}")
    wit_df = pd.read_json(args.input_path, lines=args.input_path.endswith(".jsonl"))

    logger.info(f"Reading instructions from {args.instructions_path}")
    instructions = json.load(open(args.instructions_path))

    wit_conversations_df = apply_conversations_format(
        wit_df,
        instructions=instructions,
        seed=args.seed,
    )

    logger.info(f"Saving conversations to {args.output_path}")
    wit_conversations_df.to_json(
        args.output_path, orient="records", lines=args.output_path.endswith(".jsonl")
    )
