import json
import orthanc

def GetStudyInfo(output, uri, **request):
    studyId = request['groups'][0]

    # Fetch all series and instances once
    series = json.loads(orthanc.RestApiGet(f'/studies/{studyId}/series'))
    instances = json.loads(orthanc.RestApiGet(f'/studies/{studyId}/instances?expand'))

    # Pre-index instances by series ID for fast lookup
    series_to_instances = {}
    for instance in instances:
        parent = instance['ParentSeries']
        if parent not in series_to_instances:
            series_to_instances[parent] = []
        series_to_instances[parent].append(instance)

    result = []

    for serie in series:
        tags = serie['MainDicomTags'].copy()
        tags['OrthancSeriesID'] = serie['ID']
        tags['ParentStudy'] = serie['ParentStudy']

        series_instances_info = []

        # Get instances that belong to this series
        for inst in series_to_instances.get(serie['ID'], []):
            inst_tags = inst['MainDicomTags'].copy()
            inst_tags['OrthancInstanceID'] = inst['ID']

            # Fetch metadata once per instance
            metadata = json.loads(
                orthanc.RestApiGet(f'/instances/{inst["ID"]}/metadata?expand')
            )
            inst_tags.update(metadata)

            series_instances_info.append(inst_tags)

        tags['Instances'] = series_instances_info
        result.append(tags)

    output.AnswerBuffer(json.dumps(result), 'application/json')

orthanc.RegisterRestCallback('/studies/(.*)/info', GetStudyInfo)
