import os
import io
import pandas as pd
from PIL import Image
from src.modeling.multimodal import ImageTextModelConfig, AutoImageTextGenerator, ImageTextGenerator
from src.logger import get_logger

logger = get_logger(__name__)

def build_prompt(row: pd.Series, model_name: str) -> str:
    prompt = ""
    # if "gemma" in model_name.lower():
    #     logger.info("Using Paligemma prompt format.")
    #     prompt += "answer pl "

    prompt += row["question"]
    if row["hint"]:
        prompt += "\nPodpowiedź: "
        prompt += row["hint"]

    return prompt

def read_image(image_path: str) -> Image.Image:
    if image_path.lower().endswith(".svg"):
        import cairosvg
        png_data = cairosvg.svg2png(url=image_path)
        return Image.open(io.BytesIO(png_data)).convert("RGB")

    return Image.open(image_path).convert("RGB")

def evaluate(df: pd.DataFrame, images_dir: str, model_name_or_path: str) -> pd.DataFrame:
    prompts = df.apply(lambda row: build_prompt(row, model_name_or_path), axis=1).tolist()
    images = df["image"].apply(lambda x: read_image(os.path.join(images_dir, x))).tolist()

    config = ImageTextModelConfig(model_name_or_path=model_name_or_path)
    with AutoImageTextGenerator.from_config(config) as generator:
        responses = generator.generate(prompts, images)

    results_df = df.copy()
    results_df["response"] = responses

    return results_df


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file", type=str, required=True)
    parser.add_argument("--images_dir", type=str, required=True)
    parser.add_argument("--model_name_or_path", type=str, required=True)
    parser.add_argument("--output_file", type=str, required=True)
    args = parser.parse_args()

    df = pd.read_json(args.input_file, lines=args.input_file.endswith(".jsonl"))
    results_df = evaluate(df, args.images_dir, args.model_name_or_path)

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