from .imagen_museum import TASK_DICT, DOMAIN
from .imagen_museum import fetch_indexes, fetch_indexes_no_csv
import random

ARENA_TO_IG_MUSEUM = {"LCM(v1.5/XL)":"LCM", "PlayGroundV2.5": "PlayGroundV2_5", "FLUX.1-schnell": "FLUX1schnell", "FLUX.1-dev": "FLUX1dev"}
ARENA_TO_VG_MUSEUM = {"StableVideoDiffusion": "FastSVD", "CogVideoX-2B": "CogVideoX", "CogVideoX-5B": "CogVideoX5B"}

def draw2_from_imagen_museum(task, model_name1, model_name2):
    task_name = TASK_DICT[task]
    model_name1 = ARENA_TO_IG_MUSEUM[model_name1] if model_name1 in ARENA_TO_IG_MUSEUM else model_name1
    model_name2 = ARENA_TO_IG_MUSEUM[model_name2] if model_name2 in ARENA_TO_IG_MUSEUM else model_name2

    domain = DOMAIN
    baselink = domain + task_name

    matched_results = fetch_indexes(baselink)
    r = random.Random()
    uid, value = r.choice(list(matched_results.items()))
    image_link_1 = baselink + "/" + model_name1 + "/" + uid
    image_link_2 = baselink + "/" + model_name2 + "/" + uid

    if task == "t2i": # Image Gen
        prompt = value['prompt']
        return [[image_link_1, image_link_2], [prompt]]
    if task == "tie": # Image Edit
        instruction = value['instruction']
        input_caption = value['source_global_caption']
        output_caption = value['target_global_caption']
        source_image_link = baselink + "/" + "input" + "/" + uid
        return [[source_image_link, image_link_1, image_link_2], [input_caption, output_caption, instruction]]
    else:
        raise ValueError("Task not supported")

def draw_from_imagen_museum(task, model_name):
    task_name = TASK_DICT[task]
    model_name = ARENA_TO_IG_MUSEUM[model_name] if model_name in ARENA_TO_IG_MUSEUM else model_name

    domain = DOMAIN
    baselink = domain + task_name

    matched_results = fetch_indexes(baselink)
    r = random.Random()
    uid, value = r.choice(list(matched_results.items()))
    model = model_name
    image_link = baselink + "/" + model + "/" + uid
    print(image_link)

    if task == "t2i": # Image Gen
        prompt = value['prompt']
        return [image_link, prompt]
    if task == "tie": # Image Edit
        instruction = value['instruction']
        input_caption = value['source_global_caption']
        output_caption = value['target_global_caption']
        source_image_link = baselink + "/" + "input" + "/" + uid
        return [[source_image_link, image_link], [input_caption, output_caption, instruction]]
    else:
        raise ValueError("Task not supported")

def draw2_from_videogen_museum(task, model_name1, model_name2):
    domain = "https://github.com/ChromAIca/VideoGenMuseum/raw/main/Museum/"
    baselink = domain + "VideoGenHub_Text-Guided_VG"

    model_name1 = ARENA_TO_VG_MUSEUM[model_name1] if model_name1 in ARENA_TO_VG_MUSEUM else model_name1
    model_name2 = ARENA_TO_VG_MUSEUM[model_name2] if model_name2 in ARENA_TO_VG_MUSEUM else model_name2

    matched_results = fetch_indexes_no_csv(baselink)
    r = random.Random()
    uid, value = r.choice(list(matched_results.items()))
    video_link_1 = baselink + "/" + model_name1 + "/" + uid
    video_link_2 = baselink + "/" + model_name2 + "/" + uid

    if task == "t2v": # Video Gen
        prompt = value['prompt_en']
        return [[video_link_1, video_link_2], [prompt]]
    else:
        raise ValueError("Task not supported")

def draw_from_videogen_museum(task, model_name):
    domain = "https://github.com/ChromAIca/VideoGenMuseum/raw/main/Museum/"
    baselink = domain + "VideoGenHub_Text-Guided_VG"

    model_name = ARENA_TO_VG_MUSEUM[model_name] if model_name in ARENA_TO_VG_MUSEUM else model_name

    matched_results = fetch_indexes_no_csv(baselink)
    r = random.Random()
    uid, value = r.choice(list(matched_results.items()))
    model = model_name
    video_link = baselink + "/" + model + "/" + uid
    print(video_link)

    if task == "t2v": # Video Gen
        prompt = value['prompt_en']
        return [video_link, prompt]
    else:
        raise ValueError("Task not supported")