import onnxruntime as ort
import sys

def check_onnx_batch_support(model_path):
    """
    ONNX 모델을 로드하여 배치 처리 지원 여부를 확인하고,
    입력 및 출력 정보를 출력하는 함수
    """
    try:
        session = ort.InferenceSession(model_path)
        print(f"✅ 모델 로드 성공: {model_path}")

        print("\n[입력 정보]")
        for i, input_meta in enumerate(session.get_inputs()):
            shape = input_meta.shape
            print(f"  - 입력 #{i}")
            print(f"    - 이름: {input_meta.name}")
            print(f"    - 형태: {shape}")
            # 배치 차원이 문자열(동적)인지 확인
            if isinstance(shape[0], str) or shape[0] is None:
                print(f"    - ✅ 배치 처리 지원 (동적 배치 크기 '{shape[0]}')")
            else:
                print(f"    - ⚠️ 배치 처리를 지원하지 않거나, 배치 크기가 '{shape[0]}'으로 고정됨")

        print("\n[출력 정보]")
        for i, output_meta in enumerate(session.get_outputs()):
            shape = output_meta.shape
            print(f"  - 출력 #{i}")
            print(f"    - 이름: {output_meta.name}")
            print(f"    - 형태: {shape}")
            if isinstance(shape[0], str) or shape[0] is None:
                print(f"    - ✅ 배치 처리 지원 (동적 배치 크기 '{shape[0]}')")

    except Exception as e:
        print(f"❌ 오류 발생: 모델을 로드하거나 분석하는 중 문제가 발생했습니다.")
        print(f"   {model_path}: {e}")

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("사용법: python check_onnx.py <ONNX_FILE_PATH_1> [ONNX_FILE_PATH_2] ...")
    else:
        for model_path in sys.argv[1:]:
            print("-" * 50)
            check_onnx_batch_support(model_path)
            print("-" * 50)