import gradio as gr
import requests
from huggingface_hub import HfApi
import pandas as pd
from datetime import datetime
import json

# Hugging Face API 초기화
api = HfApi()

def search_models(
    search_query="",
    task_filter="all",
    language_filter="all", 
    license_filter="all",
    sort_by="downloads",
    max_results=20
):
    """Hugging Face Hub에서 모델 검색"""
    
    try:
        # 검색 파라미터 설정
        search_params = {
            "limit": max_results,
            "sort": sort_by,
            "direction": -1,  # 내림차순
        }
        
        # 태스크 필터
        if task_filter != "all":
            search_params["pipeline_tag"] = task_filter
            
        # 언어 필터  
        if language_filter != "all":
            search_params["language"] = language_filter
            
        # 라이선스 필터
        if license_filter != "all":
            search_params["license"] = license_filter
            
        # 검색 쿼리
        if search_query.strip():
            search_params["search"] = search_query
            
        # API 호출
        models = api.list_models(**search_params)
        
        # 결과 처리
        results = []
        for model in models:
            # 모델 정보 추출
            model_info = {
                "Model Name": model.modelId,
                "Task": getattr(model, 'pipeline_tag', 'Unknown'),
                "Downloads": f"{getattr(model, 'downloads', 0):,}",
                "Likes": f"{getattr(model, 'likes', 0):,}",
                "Updated": getattr(model, 'lastModified', 'Unknown'),
                "License": getattr(model, 'license', 'Unknown'),
                "Link": f"https://huggingface.co/{model.modelId}"
            }
            results.append(model_info)
            
        if not results:
            return pd.DataFrame([{"Message": "검색 결과가 없습니다. 다른 검색어나 필터를 시도해보세요."}])
            
        # DataFrame으로 변환
        df = pd.DataFrame(results)
        return df
        
    except Exception as e:
        error_df = pd.DataFrame([{"Error": f"검색 중 오류가 발생했습니다: {str(e)}"}])
        return error_df

def get_model_details(model_name):
    """특정 모델의 상세 정보 가져오기"""
    if not model_name:
        return "모델명을 입력해주세요."
        
    try:
        # 모델 정보 가져오기
        model_info = api.model_info(model_name)
        
        details = f"""
# 🤗 {model_name}

## 기본 정보
- **태스크**: {getattr(model_info, 'pipeline_tag', 'Unknown')}
- **라이선스**: {getattr(model_info, 'license', 'Unknown')}
- **다운로드**: {getattr(model_info, 'downloads', 0):,}회
- **좋아요**: {getattr(model_info, 'likes', 0):,}개
- **최종 업데이트**: {getattr(model_info, 'lastModified', 'Unknown')}

## 지원 언어
{', '.join(getattr(model_info, 'language', ['Unknown']))}

## 태그
{', '.join(getattr(model_info, 'tags', ['None']))}

## 모델 링크
🔗 [Hugging Face에서 보기](https://huggingface.co/{model_name})

## 사용 예제
```python
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("{model_name}")
model = AutoModel.from_pretrained("{model_name}")
```
"""
        return details
        
    except Exception as e:
        return f"❌ 오류: {str(e)}\n\n모델명을 정확히 입력했는지 확인해주세요."

# Gradio 인터페이스
with gr.Blocks(title="🔍 HF Model Finder", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("# 🔍 Hugging Face Model Finder")
    gr.Markdown("허깅페이스의 수많은 모델 중에서 당신이 원하는 조건에 맞는 모델을 쉽게 찾아보세요!")
    
    with gr.Tabs():
        # 탭 1: 모델 검색
        with gr.Tab("🔍 모델 검색"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### 검색 조건")
                    
                    search_query = gr.Textbox(
                        label="검색어",
                        placeholder="예: korean, sentiment, bert",
                        value=""
                    )
                    
                    task_filter = gr.Dropdown(
                        label="태스크",
                        choices=[
                            "all", "text-classification", "token-classification", 
                            "question-answering", "fill-mask", "summarization",
                            "translation", "text-generation", "conversational"
                        ],
                        value="all"
                    )
                    
                    language_filter = gr.Dropdown(
                        label="언어",
                        choices=["all", "ko", "en", "multilingual", "zh", "ja", "es", "fr"],
                        value="all"
                    )
                    
                    license_filter = gr.Dropdown(
                        label="라이선스",
                        choices=["all", "apache-2.0", "mit", "cc-by-4.0", "cc-by-nc-4.0"],
                        value="all"
                    )
                    
                    sort_by = gr.Dropdown(
                        label="정렬 기준",
                        choices=["downloads", "likes", "lastModified"],
                        value="downloads"
                    )
                    
                    max_results = gr.Slider(
                        label="최대 결과 수",
                        minimum=5,
                        maximum=50,
                        value=20,
                        step=5
                    )
                    
                    search_btn = gr.Button("🔍 검색하기", variant="primary")
                
                with gr.Column(scale=2):
                    gr.Markdown("### 검색 결과")
                    results_table = gr.Dataframe(
                        headers=["Model Name", "Task", "Downloads", "Likes", "Updated", "License"],
                        interactive=False,
                        wrap=True
                    )
        
        # 탭 2: 모델 상세 정보
        with gr.Tab("📋 모델 상세 정보"):
            with gr.Row():
                with gr.Column(scale=1):
                    model_name_input = gr.Textbox(
                        label="모델명",
                        placeholder="예: klue/bert-base",
                        info="Hugging Face의 정확한 모델명을 입력하세요"
                    )
                    detail_btn = gr.Button("📋 상세 정보 보기", variant="primary")
                
                with gr.Column(scale=2):
                    model_details = gr.Markdown("모델명을 입력하고 버튼을 클릭하세요.")
    
    # 이벤트 핸들러
    search_btn.click(
        fn=search_models,
        inputs=[search_query, task_filter, language_filter, license_filter, sort_by, max_results],
        outputs=results_table
    )
    
    detail_btn.click(
        fn=get_model_details,
        inputs=model_name_input,
        outputs=model_details
    )
    
    # 사용법 안내
    gr.Markdown("""
    ### 💡 사용 팁
    
    **검색 탭:**
    - 검색어에 키워드를 입력하거나 필터를 조합해서 사용하세요
    - 다운로드 수가 많은 모델일수록 안정적이고 문서화가 잘 되어 있어요
    - 라이선스를 확인해서 상업적 사용 가능 여부를 체크하세요
    
    **상세 정보 탭:**
    - 검색 결과에서 관심 있는 모델명을 복사해서 붙여넣으세요
    - 사용 예제 코드도 자동으로 생성됩니다
    
    **추천 검색 조합:**
    - 한국어 감정분석: language=ko, task=text-classification, 검색어=sentiment
    - 영어 질의응답: language=en, task=question-answering
    - 상업적 사용 가능: license=apache-2.0 또는 mit
    """)

if __name__ == "__main__":
    demo.launch()