import os
import json
from typing import Tuple, Optional
from transformers import AutoTokenizer
from huggingface_hub import snapshot_download

def download_model(config_path: str = './config.json') -> Tuple[Optional[object], Optional[AutoTokenizer]]:
    """
    Hugging Face에서 모델 아티팩트만 안전하게 다운로드합니다.
    - 모델 로딩(메모리에 적재) 없이 파일만 받도록 snapshot_download 사용
    - Colab에서 대형 모델의 초기 로딩 문제(양자화/가속기 버전 차이 등)를 회피

    반환: (model, tokenizer)
      - 이 함수는 모델을 메모리에 로드하지 않으므로 model은 항상 None을 반환합니다.
      - tokenizer는 로컬 경로에서 로드에 성공하면 반환, 실패 시 None
    """
    with open(config_path, 'r') as f:
        config = json.load(f)

    model_name = config['model_name']
    local_path = config['model_local_path']

    os.makedirs(local_path, exist_ok=True)
    print(f"모델 {model_name}을 {local_path}에 다운로드 중...")

    try:
        # 인증 토큰(필요 시) 지원: 환경변수 HF_TOKEN 사용
        hf_token = os.environ.get('HF_TOKEN', None)

        # 전체 리포 스냅샷을 로컬 디렉토리에 동기화
        snapshot_download(
            repo_id=model_name,
            local_dir=local_path,
            local_dir_use_symlinks=False,
            resume_download=True,
            token=hf_token
        )

        # 토크나이저 로드 가능 여부만 확인 (모델은 나중에 로드)
        tokenizer = None
        try:
            tokenizer = AutoTokenizer.from_pretrained(local_path)
        except Exception:
            # 토크나이저 파일이 없을 수도 있으므로 경고만 출력
            print("토크나이저 확인 실패(계속 진행): 로컬 경로에 tokenizer 파일이 없을 수 있습니다.")

        print(f"모델 다운로드 완료: {local_path}")
        return None, tokenizer
    except Exception as e:
        print(f"모델 다운로드 실패: {e}")
        return None, None

if __name__ == "__main__":
    download_model()