# python3 main.py --endpoint gemma3
import argparse
import os
import time

import requests

# API 키
API_KEY = "sk-dade0cb396c744ec431357cedd5784c2"

# 지원하는 파일 확장자
SUPPORTED_EXTENSIONS = (".pdf", ".png", ".jpg", ".jpeg", ".docx", ".json")


def send_request(api_endpoint, file_path, model_name, prompt_file_path, source_dir):
    """지정된 파일과 옵션으로 API에 멀티모달 추출 요청을 보냅니다."""
    print(
        f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] '{os.path.basename(file_path)}' 파일 전송 중..."
    )

    try:
        headers = {"x-api-key": API_KEY}

        with open(file_path, "rb") as input_file:
            files = {"input_file": (os.path.basename(file_path), input_file)}
            data = {"source_dir": source_dir}

            if model_name:
                data["model"] = model_name

            if prompt_file_path:
                try:
                    # 프롬프트 파일을 열어 files 딕셔너리에 추가
                    prompt_file_opened = open(prompt_file_path, "rb")
                    files["prompt_file"] = (
                        os.path.basename(prompt_file_path),
                        prompt_file_opened,
                    )
                except FileNotFoundError:
                    print(
                        f"  [오류] 프롬프트 파일을 찾을 수 없습니다: {prompt_file_path}"
                    )
                    return

            response = requests.post(
                api_endpoint, headers=headers, files=files, data=data, timeout=30
            )

            # 프롬프트 파일이 열려있다면 닫기
            if "prompt_file" in files and not files["prompt_file"][1].closed:
                files["prompt_file"][1].close()

            response.raise_for_status()  # HTTP 오류 발생 시 예외 발생

            result = response.json()
            request_id = result.get("request_id")
            print(f"  [성공] 요청 ID: {request_id}")

    except requests.exceptions.RequestException as e:
        print(f"  [오류] API 요청 중 문제가 발생했습니다: {e}")
    except Exception as e:
        print(f"  [오류] 알 수 없는 오류가 발생했습니다: {e}")


def main(model_name, prompt_file_path, endpoint_name):
    """데이터 디렉토리를 순회하며 지원하는 각 파일에 대해 API 요청을 보냅니다."""
    api_endpoint = f"http://localhost:8889/costs/{endpoint_name}"
    base_data_directory = "ocr_results"
    target_dirs = ["pp-ocr", "pp-structure", "upstage"]
    total_file_count = 0
    skipped_file_count = 0

    for dir_name in target_dirs:
        data_directory = os.path.join(base_data_directory, dir_name)
        # Correctly construct the result directory path relative to the project root
        result_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'result', f"{endpoint_name}-{dir_name}")
        
        print(f"대상 디렉토리: {data_directory}")
        print(f"결과 확인 디렉토리: {result_dir}")
        print("-" * 30)

        if not os.path.isdir(data_directory):
            print(f"[오류] 디렉토리를 찾을 수 없습니다: {data_directory}")
            continue

        file_count = 0
        for root, _, files in os.walk(data_directory):
            for file in files:
                if file.lower().endswith(SUPPORTED_EXTENSIONS):
                    # 결과 파일 존재 여부 확인
                    result_file_path = os.path.join(result_dir, file)
                    if os.path.exists(result_file_path):
                        print(f"  [건너뛰기] 이미 결과가 존재합니다: {file}")
                        skipped_file_count += 1
                        continue

                    file_path = os.path.join(root, file)
                    send_request(api_endpoint, file_path, model_name, prompt_file_path, dir_name)
                    print("  [대기] 다음 요청까지 30초 대기...")
                    time.sleep(30)
                    file_count += 1
        
        total_file_count += file_count
        print(f"'{dir_name}' 디렉토리의 {file_count}개 파일 요청 완료.")

    print("-" * 30)
    print(f"총 {total_file_count}개의 파일에 대한 요청을 완료했으며, {skipped_file_count}개를 건너뛰었습니다.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="멀티모달 API에 파일들을 순차적으로 요청하는 스크립트"
    )

    parser.add_argument(
        "--endpoint",
        type=str,
        default="gemini",
        help="호출할 API 엔드포인트 이름 (예: gemini, gemma3)",
        dest="endpoint_name"
    )
    parser.add_argument(
        "--model", type=str, help="사용할 LLM 모델 이름"
    )
    parser.add_argument(
        "--prompt_file",
        type=str,
        help="구조화에 사용할 커스텀 .txt 프롬프트 파일의 경로",
    )

    args = parser.parse_args()
    main(args.model, args.prompt_file, args.endpoint_name)