import gradio as gr
import earthview as ev
import os
import json
import utils
from pandas import DataFrame

# --- Configuration ---
DATASET_SUBSET = "satellogic"
LABELED_DATA_FILE = "labeled_data.json"
SAMPLE_SEED = 10  # The seed to use when sampling the dataset for the demo.

# --- Load Dataset ---
dataset = ev.load_dataset(DATASET_SUBSET, shards=[SAMPLE_SEED])
data_iter = iter(dataset)

# --- Load Labeled Data (if it exists) ---
def load_labeled_data():
    if os.path.exists(LABELED_DATA_FILE):
        with open(LABELED_DATA_FILE, "r") as f:
            return json.load(f)
    else:
        return []

labeled_data = load_labeled_data()

# --- Get Next Sample for Labeling ---
def get_next_sample():
    global data_iter
    try:
        sample = next(data_iter)
        sample = ev.item_to_images(DATASET_SUBSET, sample)
        return sample
    except StopIteration:
        print("No more samples in the dataset.")
        return None

# --- Save Labeled Data ---
def save_labeled_data(label, state):
    global labeled_data
    sample = state["sample"]
    if sample is None:
        return "No image to label", DataFrame()

    image_bytes = sample["rgb"][0].convert("RGB").tobytes()  # Convert PIL Image to bytes
    labeled_data.append({
        "image": image_bytes,
        "metadata": sample["metadata"],
        "label": label
    })

    with open(LABELED_DATA_FILE, "w") as f:
        json.dump(labeled_data, f)

    new_sample = get_next_sample()
    if new_sample is None:
        state["sample"] = None
        return "Dataset exhausted.", DataFrame()

    state["sample"] = new_sample
    new_metadata = new_sample["metadata"]
    new_metadata["map"] = f'<a href="{utils.get_google_map_link(new_sample, DATASET_SUBSET)}" target="_blank">🧭</a>'
    return "", DataFrame([new_metadata])

# --- Gradio Interface ---
def labeling_ui():
    state = gr.State({"sample": None, "subset": DATASET_SUBSET})

    with gr.Row():
        cool_button = gr.Button("Cool")
        not_cool_button = gr.Button("Not Cool")

    table = gr.DataFrame(datatype="html")

    def initialize_labeling_ui():
        sample = get_next_sample()
        if sample is None:
            return {"sample": None, "subset": DATASET_SUBSET}, DataFrame()
        metadata = sample["metadata"]
        metadata["map"] = f'<a href="{utils.get_google_map_link(sample, DATASET_SUBSET)}" target="_blank">🧭</a>'
        return {"sample": sample, "subset": DATASET_SUBSET}, DataFrame([metadata])

    initial_state, initial_metadata = initialize_labeling_ui()
    table.value = initial_metadata
    state.value = initial_state

    cool_button.click(
        fn=lambda label, state: save_labeled_data(label, state),
        inputs=[gr.Textbox(visible=False, value="cool"), state],
        outputs=[gr.Textbox(label="Debug"), table]
    )
    not_cool_button.click(
        fn=lambda label, state: save_labeled_data(label, state),
        inputs=[gr.Textbox(visible=False, value="not cool"), state],
        outputs=[gr.Textbox(label="Debug"), table]
    )

# --- Main Interface ---
with gr.Blocks() as demo:
    gr.Markdown("# TerraNomaly")
    with gr.Tabs():
        with gr.TabItem("Labeling"):
            labeling_ui()

demo.launch(debug=True)