app.py

import os
import random
import base64
import requests
from selenium import webdriver
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.by import By
from selenium.common.exceptions import WebDriverException, TimeoutException
from PIL import Image
from io import BytesIO
from datetime import datetime
import gradio as gr
from typing import Tuple
import time
from pathlib import Path  # 추가

# 스크린샷 캐시 디렉토리 설정
CACHE_DIR = Path("screenshot_cache")
CACHE_DIR.mkdir(exist_ok=True)

# 전역 변수로 스크린샷 캐시 선언
SCREENSHOT_CACHE = {}

def get_cached_screenshot(url: str) -> str:
    """캐시된 스크린샷 가져오기 또는 새로 생성"""
    cache_file = CACHE_DIR / f"{base64.b64encode(url.encode()).decode()}.png"
    
    if cache_file.exists():
        with open(cache_file, "rb") as f:
            return base64.b64encode(f.read()).decode()
            
    return take_screenshot(url)
    
def take_screenshot(url):
    """웹사이트 스크린샷 촬영 함수 (로딩 대기 시간 추가)"""
    if url in SCREENSHOT_CACHE:
        return SCREENSHOT_CACHE[url]
        
    if not url.startswith('http'):
        url = f"https://{url}"
        
    options = webdriver.ChromeOptions()
    options.add_argument('--headless')
    options.add_argument('--no-sandbox')
    options.add_argument('--disable-dev-shm-usage')
    options.add_argument('--window-size=1080,720')
    
    try:
        driver = webdriver.Chrome(options=options)
        driver.get(url)
        
        # 명시적 대기: body 요소가 로드될 때까지 대기 (최대 10초)
        try:
            WebDriverWait(driver, 10).until(
                EC.presence_of_element_located((By.TAG_NAME, "body"))
            )
        except TimeoutException:
            print(f"페이지 로딩 타임아웃: {url}")
        
        # 추가 대기 시간을 2초로 증가
        time.sleep(2)  # 1초에서 2초로 변경
        
        # JavaScript 실행 완료 대기
        driver.execute_script("return document.readyState") == "complete"
        
        # 스크린샷 촬영
        screenshot = driver.get_screenshot_as_png()
        img = Image.open(BytesIO(screenshot))
        buffered = BytesIO()
        img.save(buffered, format="PNG")
        base64_image = base64.b64encode(buffered.getvalue()).decode()
        
        # 캐시에 저장
        SCREENSHOT_CACHE[url] = base64_image
        return base64_image
        
    except WebDriverException as e:
        print(f"스크린샷 촬영 실패: {str(e)} for URL: {url}")
        return None
    except Exception as e:
        print(f"예상치 못한 오류: {str(e)} for URL: {url}")
        return None
    finally:
        if 'driver' in locals():
            driver.quit()

from datetime import datetime, timedelta

def calculate_rising_rate(created_date: str, rank: int) -> int:
    """AI Rising Rate 계산"""
    # 생성일 기준 점수 계산
    created = datetime.strptime(created_date.split('T')[0], '%Y-%m-%d')
    today = datetime.now()
    days_diff = (today - created).days
    date_score = max(0, 300 - days_diff)  # 최대 300점
    
    # 순위 기준 점수 계산
    rank_score = max(0, 300 - rank)  # 최대 300점
    
    # 총점 계산
    total_score = date_score + rank_score
    
    # 별 개수 계산 (0~5)
    if total_score <= 100:
        stars = 1
    elif total_score <= 200:
        stars = 2
    elif total_score <= 300:
        stars = 3
    elif total_score <= 400:
        stars = 4
    else:
        stars = 5
        
    return stars

def get_popularity_grade(likes: int, stars: int) -> tuple:
    """AI Popularity Score 등급 계산"""
    # 기본 점수 (likes)
    base_score = min(likes, 10000)  # 최대 10000점
    
    # 별점 추가 점수 (별 하나당 500점)
    star_score = stars * 500
    
    # 총점
    total_score = base_score + star_score
    
    # 등급 테이블 (18단계)
    grades = [
        (9000, "AAA+"), (8500, "AAA"), (8000, "AAA-"),
        (7500, "AA+"), (7000, "AA"), (6500, "AA-"),
        (6000, "A+"), (5500, "A"), (5000, "A-"),
        (4500, "BBB+"), (4000, "BBB"), (3500, "BBB-"),
        (3000, "BB+"), (2500, "BB"), (2000, "BB-"),
        (1500, "B+"), (1000, "B"), (500, "B-")
    ]
    
    for threshold, grade in grades:
        if total_score >= threshold:
            return grade, total_score
            
    return "B-", total_score

# get_card 함수 내의 hardware_info 부분을 다음으로 교체:
def get_rating_info(item: dict, index: int) -> str:
    """평가 정보 HTML 생성"""
    created = item.get('createdAt', '').split('T')[0]
    likes = int(str(item.get('likes', '0')).replace(',', ''))
    
    # AI Rising Rate 계산
    stars = calculate_rising_rate(created, index + 1)
    star_html = "★" * stars + "☆" * (5 - stars)  # 채워진 별과 빈 별 조합
    
    # AI Popularity Score 계산
    grade, score = get_popularity_grade(likes, stars)
    
    # 등급별 색상 설정
    grade_colors = {
        'AAA': '#FFD700', 'AA': '#FFA500', 'A': '#FF4500',
        'BBB': '#4169E1', 'BB': '#1E90FF', 'B': '#00BFFF'
    }
    grade_base = grade.rstrip('+-')
    grade_color = grade_colors.get(grade_base, '#666666')
    
    return f"""
    <div style='
        margin-top: 15px;
        padding: 15px;
        background: rgba(255,255,255,0.4);
        border-radius: 10px;
        font-size: 0.9em;
        box-shadow: 0 2px 10px rgba(0,0,0,0.1);'>
        <div style='
            display: grid;
            grid-template-columns: repeat(2, 1fr);
            gap: 15px;'>
            <div style='
                color: #333;
                display: flex;
                flex-direction: column;
                gap: 5px;'>
                <span style='font-weight: bold;'>AI Rising Rate:</span>
                <span style='
                    color: #FF8C00;
                    font-size: 1.4em;
                    letter-spacing: 2px;
                    text-shadow: 1px 1px 2px rgba(0,0,0,0.1);'>{star_html}</span>
            </div>
            <div style='
                color: #333;
                display: flex;
                flex-direction: column;
                gap: 5px;'>
                <span style='font-weight: bold;'>AI Popularity Score:</span>
                <span style='
                    font-size: 1.2em;
                    font-weight: bold;
                    color: {grade_color};
                    text-shadow: 1px 1px 2px rgba(0,0,0,0.1);'>{grade} ({score:,})</span>
            </div>
        </div>
    </div>
    """

def get_hardware_info(item: dict) -> tuple:
    """하드웨어 정보 추출"""
    try:
        # runtime 정보 확인
        runtime = item.get('runtime', {})
        
        # CPU 정보 처리
        cpu_info = runtime.get('cpu', 'Standard')
        
        # GPU 정보 처리
        gpu_info = "None"
        if runtime.get('accelerator') == "gpu":
            gpu_type = runtime.get('gpu', {}).get('name', '')
            gpu_memory = runtime.get('gpu', {}).get('memory', '')
            if gpu_type:
                gpu_info = f"{gpu_type}"
                if gpu_memory:
                    gpu_info += f" ({gpu_memory}GB)"
        
        # spaces decorator 확인
        if '@spaces.GPU' in str(item.get('sdk_version', '')):
            if gpu_info == "None":
                gpu_info = "GPU Enabled"
        
        # SDK 정보 처리
        sdk = item.get('sdk', 'N/A')
        
        print(f"Debug - Runtime Info: {runtime}")  # 디버그 출력
        print(f"Debug - GPU Info: {gpu_info}")    # 디버그 출력
        
        return cpu_info, gpu_info, sdk
        
    except Exception as e:
        print(f"Error parsing hardware info: {str(e)}")
        return 'Standard', 'None', 'N/A'

def get_card(item: dict, index: int, card_type: str = "space") -> str:
    """통합 카드 HTML 생성"""
    item_id = item.get('id', '')
    author, title = item_id.split('/', 1)
    likes = format(item.get('likes', 0), ',')
    created = item.get('createdAt', '').split('T')[0]
    
    # URL 정의
    if card_type == "space":
        url = f"https://huggingface.co/spaces/{item_id}"
    elif card_type == "model":
        url = f"https://huggingface.co/{item_id}"
    else:  # dataset
        url = f"https://huggingface.co/datasets/{item_id}"
    
    # 메타데이터 처리
    tags = item.get('tags', [])
    pipeline_tag = item.get('pipeline_tag', '')
    license = item.get('license', '')
    sdk = item.get('sdk', 'N/A')

    # AI Rating 정보 가져오기
    rating_info = get_rating_info(item, index)

    # 카드 타입별 그라데이션 설정
    if card_type == "space":
        gradient_colors = """
            rgba(255, 182, 193, 0.7),    /* 파스텔 핑크 */
            rgba(173, 216, 230, 0.7),    /* 파스텔 블루 */
            rgba(255, 218, 185, 0.7)     /* 파스텔 피치 */
        """
        bg_content = f"""
            background-image: url(data:image/png;base64,{get_cached_screenshot(url) if get_cached_screenshot(url) else ''});
            background-size: cover;
            background-position: center;
        """
        type_icon = "🎯"
        type_label = "SPACE"
    elif card_type == "model":
        gradient_colors = """
            rgba(110, 142, 251, 0.7),    /* 모델 블루 */
            rgba(130, 158, 251, 0.7),
            rgba(150, 174, 251, 0.7)
        """
        bg_content = f"""
            background: linear-gradient(135deg, #6e8efb, #4a6cf7);
            padding: 15px;
        """
        type_icon = "🤖"
        type_label = "MODEL"
    else:  # dataset
        gradient_colors = """
            rgba(255, 107, 107, 0.7),    /* 데이터셋 레드 */
            rgba(255, 127, 127, 0.7),
            rgba(255, 147, 147, 0.7)
        """
        bg_content = f"""
            background: linear-gradient(135deg, #ff6b6b, #ff8787);
            padding: 15px;
        """
        type_icon = "📊"
        type_label = "DATASET"

    content_bg = f"""
        background: linear-gradient(135deg, {gradient_colors});
        backdrop-filter: blur(10px);
    """


    # 태그 표시 (models와 datasets용)
    tags_html = ""
    if card_type != "space":
        tags_html = f"""
        <div style='
            position: absolute;
            top: 50%;
            left: 50%;
            transform: translate(-50%, -50%);
            display: flex;
            flex-wrap: wrap;
            gap: 5px;
            justify-content: center;
            width: 90%;'>
            {' '.join([f'''
                <span style='
                    background: rgba(255,255,255,0.2);
                    padding: 5px 10px;
                    border-radius: 15px;
                    color: white;
                    font-size: 0.8em;'>
                    #{tag}
                </span>
            ''' for tag in tags[:5]])}
        </div>
        """

    # 카드 HTML 반환
    return f"""
    <div class="card" style='
        position: relative;
        border: none; 
        padding: 0;
        margin: 10px;
        border-radius: 20px;
        box-shadow: 0 10px 20px rgba(0,0,0,0.1);
        background: white;
        transition: all 0.3s ease;
        overflow: hidden;
        min-height: 400px;
        cursor: pointer;
        transform-origin: center;'
        onmouseover="this.style.transform='scale(0.98) translateY(5px)'; this.style.boxShadow='0 5px 15px rgba(0,0,0,0.2)';"
        onmouseout="this.style.transform='scale(1) translateY(0)'; this.style.boxShadow='0 10px 20px rgba(0,0,0,0.1)';"
        onclick="window.open('{url}', '_blank')">
        
        <!-- 상단 영역 -->
        <div style='
            width: 100%;
            height: 200px;
            {bg_content}
            position: relative;'>
            
            <!-- 순위 뱃지 -->
            <div style='
                position: absolute;
                top: 10px;
                left: 10px;
                background: rgba(0,0,0,0.7);
                color: white;
                padding: 5px 15px;
                border-radius: 20px;
                font-weight: bold;
                font-size: 0.9em;
                backdrop-filter: blur(5px);'>
                #{index + 1}
            </div>
            
            <!-- 타입 뱃지 -->
            <div style='
                position: absolute;
                top: 10px;
                right: 10px;
                background: rgba(255,255,255,0.9);
                padding: 5px 15px;
                border-radius: 20px;
                font-weight: bold;
                font-size: 0.8em;'>
                {type_icon} {type_label}
            </div>

            {tags_html}
        </div>
        
        <!-- 콘텐츠 영역 -->
        <div style='
            padding: 20px;
            {content_bg}
            border-radius: 0 0 20px 20px;
            border-top: 1px solid rgba(255,255,255,0.5);'>
            <h3 style='
                margin: 0 0 15px 0;
                color: #333;
                font-size: 1.3em;
                line-height: 1.4;
                display: -webkit-box;
                -webkit-line-clamp: 2;
                -webkit-box-orient: vertical;
                overflow: hidden;
                text-overflow: ellipsis;
                text-shadow: 1px 1px 1px rgba(255,255,255,0.8);'>
                {title}
            </h3>
            
            <div style='
                display: grid;
                grid-template-columns: repeat(2, 1fr);
                gap: 10px;
                font-size: 0.9em;
                background: rgba(255,255,255,0.3);
                padding: 10px;
                border-radius: 10px;'>
                <div style='color: #444;'>
                    <span style='margin-right: 5px;'>👤</span> {author}
                </div>
                <div style='color: #444;'>
                    <span style='margin-right: 5px;'>❤️</span> {likes}
                </div>
                <div style='color: #444; grid-column: span 2;'>
                    <span style='margin-right: 5px;'>📅</span> {created}
                </div>
            </div>

            {rating_info}
        </div>
    </div>
    """

def get_trending_spaces(progress=gr.Progress()) -> Tuple[str, str]:
    """트렌딩 스페이스 가져오기"""
    url = "https://huggingface.co/api/spaces"
    
    try:
        progress(0, desc="Fetching spaces data...")
        params = {
            'full': 'true',
            'limit': 10,  # 상위 10개만 가져오도록 수정
            'sort': 'likes'  # 좋아요 순으로 정렬
        }
        response = requests.get(url, params=params)
        response.raise_for_status()
        spaces = response.json()
        
        progress(0.1, desc="Creating gallery...")
        html_content = """
        <div style='padding: 20px; background: #f5f5f5;'>
            <div style='display: grid; grid-template-columns: repeat(auto-fill, minmax(300px, 1fr)); gap: 20px;'>
        """
        
        for idx, space in enumerate(spaces):
            html_content += get_card(space, idx, "space")
            progress((0.1 + 0.9 * idx/10), desc=f"Loading space {idx+1}/10...")
            
        html_content += "</div></div>"
        
        progress(1.0, desc="Complete!")
        return html_content, "Gallery refresh complete!"
        
    except Exception as e:
        error_html = f'<div style="color: red; padding: 20px;">Error: {str(e)}</div>'
        return error_html, f"Error: {str(e)}"

def get_models(progress=gr.Progress()) -> Tuple[str, str]:
    """인기 모델 가져오기"""
    url = "https://huggingface.co/api/models"
    
    try:
        progress(0, desc="Fetching models data...")
        params = {
            'full': 'true',
            'limit': 10,  # 상위 10개만 가져오도록 수정
            'sort': 'likes'  # 좋아요 순으로 정렬
        }
        response = requests.get(url, params=params)
        response.raise_for_status()
        models = response.json()
        
        progress(0.1, desc="Creating gallery...")
        html_content = """
        <div style='padding: 20px; background: #f5f5f5;'>
            <div style='display: grid; grid-template-columns: repeat(auto-fill, minmax(300px, 1fr)); gap: 20px;'>
        """
        
        for idx, model in enumerate(models):
            html_content += get_card(model, idx, "model")
            progress((0.1 + 0.9 * idx/10), desc=f"Loading model {idx+1}/10...")
            
        html_content += "</div></div>"
        
        progress(1.0, desc="Complete!")
        return html_content, "Models gallery refresh complete!"
        
    except Exception as e:
        error_html = f'<div style="color: red; padding: 20px;">Error: {str(e)}</div>'
        return error_html, f"Error: {str(e)}"

def get_datasets(progress=gr.Progress()) -> Tuple[str, str]:
    """인기 데이터셋 가져오기"""
    url = "https://huggingface.co/api/datasets"
    
    try:
        progress(0, desc="Fetching datasets data...")
        params = {
            'full': 'true',
            'limit': 10,  # 상위 10개만 가져오도록 수정
            'sort': 'likes'  # 좋아요 순으로 정렬
        }
        response = requests.get(url, params=params)
        response.raise_for_status()
        datasets = response.json()
        
        progress(0.1, desc="Creating gallery...")
        html_content = """
        <div style='padding: 20px; background: #f5f5f5;'>
            <div style='display: grid; grid-template-columns: repeat(auto-fill, minmax(300px, 1fr)); gap: 20px;'>
        """
        
        for idx, dataset in enumerate(datasets):
            html_content += get_card(dataset, idx, "dataset")
            progress((0.1 + 0.9 * idx/10), desc=f"Loading dataset {idx+1}/10...")
            
        html_content += "</div></div>"
        
        progress(1.0, desc="Complete!")
        return html_content, "Datasets gallery refresh complete!"
        
    except Exception as e:
        error_html = f'<div style="color: red; padding: 20px;">Error: {str(e)}</div>'
        return error_html, f"Error: {str(e)}"

def create_interface():
    """Gradio 인터페이스 생성"""
    with gr.Blocks(title="HuggingFace Trending Board") as interface:
        gr.Markdown("# 🤗 HuggingFace Trending Board")
        
        with gr.Tabs() as tabs:
            # Spaces 탭
            with gr.Tab("🎯 Trending Spaces"):
                gr.Markdown("Shows top 10 trending spaces with AI ratings")
                with gr.Row():
                    spaces_refresh_btn = gr.Button("Refresh Spaces", variant="primary")
                spaces_gallery = gr.HTML()
                spaces_status = gr.Markdown("Ready")
            
            # Models 탭
            with gr.Tab("🤖 Trending Models"):
                gr.Markdown("Shows top 10 trending models with AI ratings")
                with gr.Row():
                    models_refresh_btn = gr.Button("Refresh Models", variant="primary")
                models_gallery = gr.HTML()
                models_status = gr.Markdown("Ready")
            
            # Datasets 탭
            with gr.Tab("📊 Trending Datasets"):
                gr.Markdown("Shows top 10 trending datasets with AI ratings")
                with gr.Row():
                    datasets_refresh_btn = gr.Button("Refresh Datasets", variant="primary")
                datasets_gallery = gr.HTML()
                datasets_status = gr.Markdown("Ready")
        
        # Event handlers
        spaces_refresh_btn.click(
            fn=get_trending_spaces,
            outputs=[spaces_gallery, spaces_status],
            show_progress=True
        )
        
        models_refresh_btn.click(
            fn=get_models,
            outputs=[models_gallery, models_status],
            show_progress=True
        )
        
        datasets_refresh_btn.click(
            fn=get_datasets,
            outputs=[datasets_gallery, datasets_status],
            show_progress=True
        )
        
        # 초기 로드
        interface.load(
            fn=get_trending_spaces,
            outputs=[spaces_gallery, spaces_status]
        )
        interface.load(
            fn=get_models,
            outputs=[models_gallery, models_status]
        )
        interface.load(
            fn=get_datasets,
            outputs=[datasets_gallery, datasets_status]
        )
    
    return interface

if __name__ == "__main__":
    try:
        demo = create_interface()
        demo.launch(
            share=True,
            inbrowser=True,
            show_api=False
        )
    except Exception as e:
        print(f"Error launching app: {e}")