from segments import SegmentsClient

def get_samples(client, dataset_identifier):
    page = 1
    per_page = 1000
    samples = []
    while True:
        response = client.get_samples(dataset_identifier, per_page=per_page, page=page)
        # Handle both paginated and direct list responses
        if isinstance(response, list):
            samples.extend(response)
            break  # No pagination in this case
        else:
            samples.extend(response.results)
            if not response.has_next:
                break
        page += 1
    return samples

def export_number_of_samples(samples):
    """Return the number of samples in the dataset."""
    return len(samples)

def export_frames_and_annotations(label):
    """Export frames and their annotations for non-multisensor (no sensors) labels."""
    frames = getattr(label.attributes, "frames", None)
    if frames is None:
        return []
    result = []
    for frame in frames:
        annotations = getattr(frame, "annotations", [])
        result.append({
            "frame": frame,
            "annotations": annotations
        })
    return result

def export_sensor_frames_and_annotations(label, sensor_name):
    """Export frames and annotations for a specific sensor in a multisensor label."""
    sensors = getattr(label.attributes, "sensors", None)
    if sensors is None:
        return []
    for sensor in sensors:
        if getattr(sensor, "name", None) == sensor_name:
            sensor_attrs = getattr(sensor, "attributes", None)
            if sensor_attrs is None:
                return []
            frames = getattr(sensor_attrs, "frames", None)
            if frames is None:
                return []
            result = []
            for frame in frames:
                annotations = getattr(frame, "annotations", [])
                result.append({
                    "frame": frame,
                    "annotations": annotations
                })
            return result
    return []

def export_all_sensor_frames_and_annotations(label):
    """Export all frames and annotations for all sensors in a multisensor label."""
    sensors = getattr(label.attributes, "sensors", None)
    if sensors is None:
        return []
    result = {}
    for sensor in sensors:
        sensor_name = getattr(sensor, "name", None)
        sensor_attrs = getattr(sensor, "attributes", None)
        if sensor_attrs is None:
            result[sensor_name] = []
            continue
        frames = getattr(sensor_attrs, "frames", None)
        if frames is None:
            result[sensor_name] = []
            continue
        sensor_result = []
        for frame in frames:
            annotations = getattr(frame, "annotations", [])
            sensor_result.append({
                "frame": frame,
                "annotations": annotations
            })
        result[sensor_name] = sensor_result
    return result

def main():
    api_key = DEFAULT_API_KEY
    dataset_identifier = DEFAULT_DATASET_IDENTIFIER

    client = SegmentsClient(api_key)
    samples = get_samples(client, dataset_identifier)
    number_of_samples = export_number_of_samples(samples)
    print(number_of_samples)

    first_sample = samples[0]
    label = client.get_label(first_sample.uuid)
    frames_and_annotations = export_frames_and_annotations(label)
    print(frames_and_annotations)

    sensor_name = "sensor_name"
    sensor_frames_and_annotations = export_sensor_frames_and_annotations(label, sensor_name)
    print(sensor_frames_and_annotations)

    all_sensor_frames_and_annotations = export_all_sensor_frames_and_annotations(label)
    print(all_sensor_frames_and_annotations)

if __name__ == "__main__":
    main()