import re
from typing import Callable
import pandas as pd
from tqdm import tqdm
from PIL import Image
from src.datasets.images.shards import ShardedImagesRepository
from src.logger import get_logger
from src.modeling.multimodal import ImageTextModelConfig, AutoImageTextGenerator

logger = get_logger(__name__)


def conversation_to_turns(conversation: list[dict]) -> list[tuple[dict, dict]]:
    human_messages = [msg for msg in conversation if msg["from"] == "human"]
    gpt_messages = [msg for msg in conversation if msg["from"] == "gpt"]
    assert len(human_messages) == len(gpt_messages), "Mismatched number of human and gpt messages"

    turns = zip(human_messages, gpt_messages)
    return list(turns)

def turns_to_conversation(turns: list[tuple[dict, dict]]) -> list[dict]:
    conversation = []
    for human_msg, gpt_msg in turns:
        conversation.append(human_msg)
        conversation.append(gpt_msg)
    return conversation


def create_samples(turns: list[tuple], image_store: ShardedImagesRepository, config: ImageTextModelConfig) -> tuple[list[str], list[Callable[[], Image.Image]]]:
    if not turns:
        return [], []

    samples, images = [], []
    for _, (image_source, image_name), (human_msg, gpt_msg) in turns:
        sample = config.prompt_template.format(question=human_msg["value"], answer=gpt_msg["value"])
        image = lambda s=image_source, n=image_name: image_store.get_image(s, n)
        samples.append(sample)
        images.append(image)

    return samples, images


def extract_scores(judgements: list[str]) -> list[tuple]:
    scores = []
    for judgement in judgements:
        text_quality_match = re.search(r"TextQuality:\s*(\d+)", judgement)
        visual_grounding_match = re.search(r"QuestionVisualGrounding:\s*(\d+)", judgement)
        answer_quality_match = re.search(r"AnswerQuality:\s*(\d+)", judgement)

        text_quality = int(text_quality_match.group(1)) if text_quality_match else None
        visual_grounding = int(visual_grounding_match.group(1)) if visual_grounding_match else None
        answer_quality = int(answer_quality_match.group(1)) if answer_quality_match else None
        
        scores.append((text_quality, visual_grounding, answer_quality))
    return scores


def judge(df: pd.DataFrame, image_store: ShardedImagesRepository, config: ImageTextModelConfig) -> pd.DataFrame:
    all_turns = [] # (id, image, turn)
    for _, row in df.iterrows():
        conversation = row["conversations"]
        for turn in conversation_to_turns(conversation):
            image = (row["image_source"], row["image"])
            all_turns.append((row["id"], image, turn))

    with AutoImageTextGenerator.from_config(config) as generator:
        samples, images = create_samples(all_turns, image_store, config)
        outputs = generator.generate(samples, images)

    scores = extract_scores(outputs)

    for turn, score in zip(all_turns, scores):
        id, _, (human_msg, gpt_msg) = turn
        text_quality, question_visual_grounding, answer_quality = score
        
        human_msg.setdefault("metadata", {})
        human_msg["metadata"]["question_visual_grounding"] = question_visual_grounding

        gpt_msg.setdefault("metadata", {})
        gpt_msg["metadata"]["text_quality"] = text_quality
        gpt_msg["metadata"]["answer_quality"] = answer_quality

    all_turns_grouped = {}
    for id, _, turn in all_turns:
        all_turns_grouped.setdefault(id, []).append(turn)

    judged_df = df.copy()
    judged_df["conversations"] = [turns_to_conversation(all_turns_grouped[id]) for id in judged_df["id"]]

    return judged_df
        

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

    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file", type=str, required=True, help="Path to the input file")
    parser.add_argument("--images_dir", type=str, required=True, help="Path to the images directory")
    parser.add_argument("--output_file", type=str, required=True, help="Path to the output file")
    parser.add_argument("--judge_config", type=str, required=True, help="Path to the judge configuration file")
    args = parser.parse_args()

    logger.info("Loading data from %s", args.input_file)
    df = pd.read_json(args.input_file, lines=True)
    # df = df.sample(n=100)

    logger.info("Setting up sharded images repository from %s", args.images_dir)
    image_store = ShardedImagesRepository(args.images_dir)

    logger.info("Loading judge configuration from %s", args.judge_config)
    judge_config = ImageTextModelConfig.from_yaml(args.judge_config)

    logger.info("Judging...")
    judged_df = judge(df, image_store, judge_config)

    logger.info("Saving judged data to %s", args.output_file)
    os.makedirs(os.path.dirname(args.output_file), exist_ok=True)
    use_lines = args.output_file.endswith(".jsonl")
    judged_df.to_json(args.output_file, lines=use_lines, orient="records", indent=2 if not use_lines else None)

