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


logger = get_logger(__name__)


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

    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["id"].apply(lambda x: f"{source}_{x}")
    conversations_df["image"] = df["image"]
    conversations_df["image_source"] = source
    conversations_df["conversations"] = df["value"].apply(get_conversation)

    return conversations_df


if __name__ == "__main__":
    import os
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input_path", type=str, required=True, help="Input path containing raw data."
    )
    parser.add_argument(
        "--instructions_path", type=str, required=True, help="Path to the instructions JSON file."
    )
    parser.add_argument(
        "--output_path", type=str, required=True, help="Output path to save processed data."
    )
    parser.add_argument("--source", type=str, required=True, help="Image source")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility.")

    args = parser.parse_args()

    logger.info(f"Reading synthdog data from {args.input_path}")
    synthdog_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))

    synthdog_conversations_df = apply_conversations_format(
        synthdog_df,
        instructions=instructions,
        source=args.source,
        seed=args.seed,
    )

    logger.info(f"Saving processed data to {args.output_path}")
    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
    use_lines = args.output_path.endswith(".jsonl")
    synthdog_conversations_df.to_json(
        args.output_path, orient="records", lines=use_lines, indent=2 if not use_lines else None
    )
