#!/usr/bin/env python3
"""
Hugging Face 클라우드 GPU 환경 설정
"""

import os
import requests
import json
import logging
from huggingface_hub import login, HfApi
from transformers import AutoTokenizer, AutoModelForCausalLM

# 로깅 설정
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class HuggingFaceCloudSetup:
    """Hugging Face 클라우드 GPU 설정 클래스"""
    
    def __init__(self):
        """초기화"""
        self.api = None
        self.model_name = "heegyu/polyglot-ko-5.8b-chat"
        self.space_name = "lily-math-rag"
        
    def setup_huggingface_login(self):
        """Hugging Face 로그인 설정"""
        logger.info("🔐 Hugging Face 로그인 설정")
        
        try:
            # 토큰 입력 요청
            token = input("Hugging Face 토큰을 입력하세요: ").strip()
            
            if token:
                login(token)
                self.api = HfApi(token=token)
                logger.info("✅ Hugging Face 로그인 성공")
                return True
            else:
                logger.error("❌ 토큰이 필요합니다")
                return False
                
        except Exception as e:
            logger.error(f"❌ Hugging Face 로그인 실패: {e}")
            return False
    
    def create_inference_endpoint(self):
        """추론 엔드포인트 생성"""
        logger.info("🚀 추론 엔드포인트 생성 중...")
        
        try:
            # 엔드포인트 설정
            endpoint_config = {
                "account": "your-username",  # Hugging Face 사용자명
                "name": "lily-math-rag-endpoint",
                "repository": self.model_name,
                "framework": "pytorch",
                "accelerator": "gpu",
                "instance_type": "gpu.t4.medium",  # GPU 인스턴스 타입
                "region": "us-east-1",
                "vendor": "aws"
            }
            
            logger.info("✅ 엔드포인트 설정 완료")
            logger.info(f"   모델: {self.model_name}")
            logger.info(f"   GPU: {endpoint_config['instance_type']}")
            logger.info(f"   지역: {endpoint_config['region']}")
            
            return endpoint_config
            
        except Exception as e:
            logger.error(f"❌ 엔드포인트 생성 실패: {e}")
            return None
    
    def create_huggingface_space(self):
        """Hugging Face Space 생성"""
        logger.info("🌐 Hugging Face Space 생성 중...")
        
        try:
            # Space 설정
            space_config = {
                "name": self.space_name,
                "type": "gradio",
                "sdk": "gradio",
                "title": "Lily Math RAG System",
                "description": "수학 문제 해결을 위한 RAG 시스템",
                "license": "mit",
                "python_version": "3.9"
            }
            
            logger.info("✅ Space 설정 완료")
            logger.info(f"   Space 이름: {space_config['name']}")
            logger.info(f"   타입: {space_config['type']}")
            
            return space_config
            
        except Exception as e:
            logger.error(f"❌ Space 생성 실패: {e}")
            return None
    
    def upload_model_to_hub(self):
        """모델을 Hugging Face Hub에 업로드"""
        logger.info("📤 모델 업로드 중...")
        
        try:
            # 로컬 모델 경로
            local_model_path = "hearth_llm_model"
            
            if os.path.exists(local_model_path):
                logger.info(f"✅ 로컬 모델 발견: {local_model_path}")
                
                # 모델 업로드 (실제로는 Hugging Face CLI 사용)
                logger.info("💡 다음 명령어로 모델을 업로드하세요:")
                logger.info(f"   huggingface-cli upload your-username/lily-math-model {local_model_path}")
                
                return True
            else:
                logger.warning(f"⚠️ 로컬 모델을 찾을 수 없습니다: {local_model_path}")
                return False
                
        except Exception as e:
            logger.error(f"❌ 모델 업로드 실패: {e}")
            return False
    
    def test_cloud_inference(self, endpoint_url):
        """클라우드 추론 테스트"""
        logger.info("🧪 클라우드 추론 테스트")
        
        try:
            # 테스트 요청
            test_data = {
                "inputs": "안녕하세요! 수학 문제를 도와주세요.",
                "parameters": {
                    "max_length": 100,
                    "temperature": 0.7
                }
            }
            
            response = requests.post(
                f"{endpoint_url}/predict",
                json=test_data,
                headers={"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
            )
            
            if response.status_code == 200:
                result = response.json()
                logger.info(f"✅ 추론 테스트 성공: {result}")
                return True
            else:
                logger.error(f"❌ 추론 테스트 실패: {response.status_code}")
                return False
                
        except Exception as e:
            logger.error(f"❌ 추론 테스트 실패: {e}")
            return False

def main():
    """메인 설정 함수"""
    print("🚀 Hugging Face 클라우드 GPU 환경 설정")
    print("=" * 50)
    
    # 1. Hugging Face 설정 클래스 초기화
    setup = HuggingFaceCloudSetup()
    
    # 2. Hugging Face 로그인
    if not setup.setup_huggingface_login():
        print("❌ 로그인 실패")
        return
    
    # 3. 추론 엔드포인트 생성
    endpoint_config = setup.create_inference_endpoint()
    
    # 4. Hugging Face Space 생성
    space_config = setup.create_huggingface_space()
    
    # 5. 모델 업로드
    setup.upload_model_to_hub()
    
    print("\n🎉 Hugging Face 클라우드 설정 완료!")
    print("✅ 다음 단계:")
    print("1. Hugging Face Inference Endpoints에서 엔드포인트 생성")
    print("2. 모델을 Hugging Face Hub에 업로드")
    print("3. Railway Hearth Chat과 연동 설정")

if __name__ == "__main__":
    main()