import json
import os
import random
from typing import Optional
import pandas as pd
from src.logger import get_logger

logger = get_logger(__name__)


def apply_conversations_format(df: pd.DataFrame, lang: str = "pl", random_state: int = 42) -> pd.DataFrame:
    random.seed(random_state)

    logger.info("Applying conversations format")

    def get_id(row: pd.Series) -> str:
        return "synthdog_" + str(row["id"])

    def get_conversations(text: str, image_token: str = "<image>") -> list[dict[str, str]]:
        instruction = random.choice(instructions[lang])
        image_position = random.choice(["preceding", "following"])
        return [
            {
                "from": "human",
                "value": (
                    (f"{image_token}\n" if image_position == "preceding" else "") +
                    instruction +
                    (f"\n{image_token}" if image_position == "following" else "")
                ),
            },
            {"from": "gpt", "value": text},
        ]

    conversations_df = df[["id", "image"]].copy()
    conversations_df["id"] = df.apply(get_id, axis=1)
    conversations_df["conversations"] = df["value"].apply(get_conversations)

    return conversations_df


def apply_images_source(df: pd.DataFrame, source: Optional[str] = None) -> pd.DataFrame:
    logger.info(f"Applying images source: {source}")
    df["image_source"] = source
    return df


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input_path", type=str, required=True, help="Input path containing raw data."
    )
    parser.add_argument(
        "--output_path", type=str, required=True, help="Output path to save processed data."
    )
    parser.add_argument(
        "--lang", type=str, default="pl", help="Language of the dataset."
    )
    parser.add_argument(
        "--image_source", type=str, default="synthdog_pl_500k", help="Source identifier for images."
    )
    parser.add_argument(
        "--random_state", type=int, default=42, help="Random state for reproducibility."
    )

    args = parser.parse_args()

    df = pd.read_json(args.input_path, lines=args.input_path.endswith(".jsonl"))

    conversations_df = apply_conversations_format(df, lang=args.lang, random_state=args.random_state)
    conversations_df = apply_images_source(conversations_df, source=args.image_source)

    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
    use_jsonl = args.output_path.endswith(".jsonl")
    conversations_df.to_json(args.output_path, orient="records", lines=use_jsonl)
    logger.info(f"Saved processed data to {args.output_path}")