import os
import pandas as pd
from src.logger import get_logger
from src.modeling.llava.predict import LlavaNextModelConfig, LlavaNextGenerator
from PIL import Image

logger = get_logger(__name__)


def evaluate(
    model_name_or_path: str,
    mllm_guard_df: pd.DataFrame,
    images_dir: str,
    lang: str = "pl",
) -> pd.DataFrame:
    logger.info(f"Loading model from {model_name_or_path}")

    model_config = LlavaNextModelConfig(model_name_or_path=model_name_or_path)
    model = LlavaNextGenerator(model_config)
    logger.info("Model loaded successfully")

    def build_prompt(row):
        prompt = row[f"Prompt_{lang}"]
        options = row[f"Options_{lang}"]

        if options:
            prompt += "\nOdpowiedzi:\n"
            prompt += options

        return prompt

    def get_image(row):
        image_path = os.path.join(images_dir, row["Image_path"])
        return Image.open(image_path).convert("RGB")

    prompts = mllm_guard_df.apply(build_prompt, axis=1).tolist()
    images = mllm_guard_df.apply(get_image, axis=1).tolist()

    logger.info("Generating results...")
    results = model.generate(prompts, images)

    mllm_guard_df[f"Response_{lang}"] = results
    return mllm_guard_df


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file", type=str, required=True, help="Input file path")
    parser.add_argument(
        "--images_dir", type=str, required=True, help="Directory with images"
    )
    parser.add_argument(
        "--model_name_or_path", type=str, required=True, help="Model name or path"
    )
    parser.add_argument(
        "--output_file", type=str, required=True, help="Output file path"
    )
    parser.add_argument(
        "--lang", type=str, default="pl", help="Language for prompts and responses"
    )

    args = parser.parse_args()

    mllm_guard_df = pd.read_json(args.input_file)
    result_df = evaluate(args.model_name_or_path, mllm_guard_df, args.images_dir, args.lang)

    os.makedirs(os.path.dirname(args.output_file), exist_ok=True)
    use_lines = args.output_file.endswith(".jsonl")
    result_df.to_json(
        args.output_file,
        orient="records",
        lines=use_lines,
        indent=2 if not use_lines else None,
    )
