from typing import Optional
import pandas as pd
from src.logger import get_logger
from src.generator import ModelConfig, LLMGenerator


logger = get_logger(__name__)


def translate_wrapped(data: pd.DataFrame, config: ModelConfig) -> pd.DataFrame:
    """
    Translates the LLaVA dataset conversations. Assumes the conversations contain a "en" key for English text.
    """
    logger.info("Setting up the translator with the provided configuration.")
    logger.info(f"Using the following configuration: {repr(config)}")
    translator = LLMGenerator(config)

    logger.info("Starting translation of conversations.")

    prompt = lambda text: config.prompt_template.format(text=text).strip()

    conversations = data["conversations"].tolist()
    items = [
        (index, item)
        for index, conversation in enumerate(conversations)
        for item in conversation
        if "en" in item
    ]
    items_translated = []

    samples = [item["en"] for _, item in items]
    samples_translated = translator.generate(samples, prompt)

    for (index, item), sample_translated in zip(items, samples_translated):
        item_translated = item.copy()
        item_translated["pl"] = sample_translated
        items_translated.append((index, item_translated))

    conversations_translated = [[] for _ in conversations]
    for index, item_translated in items_translated:
        conversations_translated[index].append(item_translated)

    translated_data = data.copy()
    translated_data["conversations"] = conversations_translated

    logger.info("Translation completed successfully.")
    return translated_data


def translate_unwrapped(data: pd.DataFrame, config: ModelConfig, mask: Optional[pd.Series] = None) -> pd.DataFrame:
    """
    Translates the LLaVA dataset conversations that have been previously unwrapped.
    """
    logger.info("Setting up the translator with the provided configuration.")
    logger.info(f"Using the following configuration: {repr(config)}")
    translator = LLMGenerator(config)

    logger.info("Starting translation of conversations.")

    prompt = lambda text: config.prompt_template.format(text=text).strip()

    samples = data.loc[mask, "en"].tolist() if mask is not None else data["en"].tolist()
    samples_translated = translator.generate(samples, prompt)
    
    translated_data = data.copy()
    if mask is not None:
        translated_data['pl'] = pd.NA
        translated_data.loc[mask, "pl"] = samples_translated
    else:
        translated_data["pl"] = samples_translated

    logger.info("Translation completed successfully.")
    return translated_data


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=str, help="Path to the input file", required=True)
    parser.add_argument("--config", type=str, help="Path to the configuration file", required=True)
    parser.add_argument("--output", type=str, help="Path to the output file", required=True)
    parser.add_argument("--num_gpus", type=int, default=1, help="Number of GPUs to use for translation")
    parser.add_argument("--unwrapped", action="store_true", help="Indicates if the input data is wrapped or unwrapped")
    parser.add_argument("--translate_instructions", action="store_true", help="Translate instructions in the datasetk")

    args = parser.parse_args()

    data = pd.read_json(args.input, lines=args.input.endswith(".jsonl"))
    config = ModelConfig.from_yaml(args.config)

    config.parallel_size = args.num_gpus

    if args.unwrapped:
        logger.info("Translating unwrapped data.")
        mask = data["from"] == "gpt" if not args.translate_instructions else None
        translated_data = translate_unwrapped(data, config, mask)
    else:
        logger.info("Translating wrapped data.")
        translated_data = translate_wrapped(data, config)

    logger.info("Saving the translated data to the output file.")
    translated_data.to_json(args.output, orient="records", lines=args.output.endswith(".jsonl"))
