from typing import Literal
import pandas as pd


def get_translated_conversation(
    conversation: list[dict], lang: str = "pl"
) -> list[dict]:
    image_token = "<image>"
    translated_conversation = []
    for item in conversation:
        if item["from"] == "human":
            trasnlated_item = {
                "from": "human",
                "value": (
                    (
                        image_token + "\n"
                        if item["image_position"] == "preceding"
                        else ""
                    )
                    + item[lang]
                    + (
                        "\n" + image_token
                        if item["image_position"] == "following"
                        else ""
                    )
                ),
            }
        elif item["from"] == "gpt":
            trasnlated_item = {"from": "gpt", "value": item[lang]}

        translated_conversation.append(trasnlated_item)

    return translated_conversation


def word_count(text: str) -> int:
    return len(text.split())


def extract_image(
    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: extract_image(x["value"])
        )
    )
    unwrapped["en"] = raw_values
    unwrapped["image_position"] = image_positions

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

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

    assert lang in df.columns, f"Language column '{lang}' not found in DataFrame columns."
    drop_cols = {'index', 'from', 'image_position', 'conversations', 'en', lang}

    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[lang]
                    + ("\n" + image_token if row['image_position'] == 'following' else "")
                )
            }
            conversation.append(item)
        conversations.append(conversation)

    wrapped_df['conversations'] = conversations

    return wrapped_df.reset_index(drop=True)
