from functools import lru_cache
import pandas as pd
import evaluate
import requests
from PIL import Image
from io import BytesIO
from typing import Literal
from tqdm import tqdm


def word_count(text: str) -> int:
    if not isinstance(text, str) or not text.strip():
        return 0

    return len(text.split())


def image_from_url(url):
    response = requests.get(url)
    response.raise_for_status()

    img = Image.open(BytesIO(response.content))
    return img


def add_image_token(
    value: str, position: Literal["preceding", "following"], image_token: str = "<image>"
) -> str:
    if position == "preceding":
        return f"{image_token}\n{value}"
    elif position == "following":
        return f"{value}\n{image_token}"

    raise ValueError(f"Invalid position: {position}. Must be 'preceding' or 'following'.")


def instruction_image_split(
    value: str, image_token: str = "<image>"
) -> tuple[str, Literal["preceding", "following", None]]:
    im_idx = value.find(image_token)
    if im_idx == -1:
        return value, None

    if im_idx == 0:
        im_pos = "preceding"
    else:
        im_pos = "following"

    raw = value.replace(image_token, "").strip()
    return raw, im_pos


def unwrap(raw_df: pd.DataFrame) -> pd.DataFrame:
    """
    Unwraps the DataFrame by expanding the 'conversations' column into multiple rows.
    """
    df = raw_df.copy()
    unwrapped = df.explode("conversations", ignore_index=False)
    unwrapped["index"] = unwrapped.groupby(level=0).cumcount()
    unwrapped["from"] = unwrapped["conversations"].apply(lambda x: x["from"])

    raw_values, image_positions = zip(
        *unwrapped["conversations"].apply(lambda x: instruction_image_split(x["value"]))
    )
    unwrapped["value"] = raw_values
    unwrapped["image_position"] = image_positions
    unwrapped["metadata"] = unwrapped["conversations"].apply(lambda x: x.get("metadata", {}))

    return unwrapped.drop(columns=["conversations"]).reset_index(drop=True)


def wrap(unwrapped_df: pd.DataFrame, image_token: str = "<image>") -> pd.DataFrame:
    """
    Wraps the DataFrame by combining rows into a single 'conversations' column.
    """
    df = unwrapped_df.copy()

    wrapped_group = df.groupby("id")

    wrapped_df = pd.DataFrame()
    wrapped_df["id"] = wrapped_group["id"].first()
    wrapped_df["image"] = wrapped_group["image"].first()

    conversations = []
    for _, group in wrapped_group:
        group = group.sort_values("index")
        conversation = []
        for _, row in group.iterrows():
            item = {
                "from": row["from"],
                "value": (
                    (image_token + "\n" if row["image_position"] == "preceding" else "")
                    + row["value"]
                    + ("\n" + image_token if row["image_position"] == "following" else "")
                ),
                "metadata": row.get("metadata", {}),
            }
            conversation.append(item)
        conversations.append(conversation)

    wrapped_df["conversations"] = conversations

    return wrapped_df.reset_index(drop=True)


@lru_cache(maxsize=1)
def get_comet_evaluator():
    return evaluate.load("comet", "Unbabel/wmt20-comet-qe-da")


def calculate_comet_score(sources: list[str], targets: list[str], batch_size=25_000) -> list[float]:
    comet = get_comet_evaluator()
    assert len(sources) == len(targets), "Sources and targets must have the same length."

    scores = []
    for i in tqdm(range(0, len(sources), batch_size), desc="Calculating COMET scores", unit="batch"):
        batch_sources = sources[i : i + batch_size]
        batch_targets = targets[i : i + batch_size]

        results = comet.compute(
            sources=batch_sources,
            predictions=batch_targets,
            references=[""] * len(batch_sources),  # Skip references
        )

        assert results, "COMET results are empty. Check the input data."
        scores.extend(results['scores'])

    return scores
