import streamlit as st
from segments import SegmentsClient
import sys
import copy
from get_labels_from_samples import export_all_sensor_frames_and_annotations

# ---------------- Streamlit UI ----------------
st.title("Copy annotations from one dataset to another")
st.subheader("For Multi-sensor datasets - Cuboids and bounding boxes")

# Initialize success message in session state
if "success" not in st.session_state:
    st.session_state.success = ""

api_key = st.text_input("Enter your API Key", type="password")
source_uuid = st.text_input("Source UUID (the sample you want to copy from)", value="")
source_frame_num = st.number_input("Source Frame Number", min_value=1, value=50, step=1)
target_uuid = st.text_input("Target UUID (the sample you want to copy to)", value="")
target_frame_num = st.number_input("Target Frame Number", min_value=1, value=1, step=1)

if st.button("Overwrite Target Frame Annotations"):
    st.session_state.success = ""  # Clear previous success message
    if not api_key or not source_uuid or not target_uuid:
        st.error("Please fill in all fields.")
    else:
        try:
            client = SegmentsClient(api_key)
            # Fetch source and target labels
            source_label = client.get_label(source_uuid)
            target_label = client.get_label(target_uuid)

            # Extract all sensor frames/annotations
            source_sensors = export_all_sensor_frames_and_annotations(source_label)
            target_sensors = export_all_sensor_frames_and_annotations(target_label)

            # Defensive copy
            new_target_label = copy.deepcopy(target_label)
            sensors_attr = getattr(new_target_label.attributes, "sensors", None)
            if sensors_attr is None:
                st.error("Target label has no sensors.")
                st.stop()

            # For each sensor in both source and target
            for sensor_idx, sensor in enumerate(sensors_attr):
                sensor_name = getattr(sensor, "name", None)
                if not sensor_name or sensor_name not in source_sensors:
                    continue  # skip sensors not found in source
                source_frames = source_sensors[sensor_name]
                target_frames = target_sensors.get(sensor_name, [])
                # Frame indices are 1-indexed for user, 0-indexed in list
                src_idx = source_frame_num - 1
                tgt_idx = target_frame_num - 1
                if src_idx >= len(source_frames) or tgt_idx >= len(target_frames):
                    st.warning(f"Sensor '{sensor_name}': Frame index out of range. Skipped.")
                    continue
                # Overwrite target frame's annotations with source frame's annotations
                src_anns = source_frames[src_idx]["annotations"]
                tgt_frame = getattr(sensor.attributes.frames[tgt_idx], "annotations", None)
                if tgt_frame is not None:
                    sensor.attributes.frames[tgt_idx].annotations = copy.deepcopy(src_anns)
                else:
                    st.warning(f"Sensor '{sensor_name}': Could not access target frame annotations.")

            # Upload the updated label to Segments.ai (only the specified frames/annotations are changed)
            client.update_label(
                target_uuid,
                labelset="ground-truth",  # Change this if you use a different labelset name
                attributes=new_target_label.attributes.model_dump()
            )
            st.session_state.success = "Target label updated and uploaded to Segments.ai!"
        except Exception as e:
            st.error(f"Error: {e}")

# Always show the success message after the button handler
if st.session_state.success:
    st.success(st.session_state.success)