import gradio as gr
from transformers import AutoTokenizer

def count_tokens(model_name, text, hf_token=None):
    """토큰 수 계산"""
    try:
        if not model_name or not text:
            return "모델명과 텍스트를 모두 입력해주세요."
        
        # 토크나이저 로드
        tokenizer = AutoTokenizer.from_pretrained(
            model_name, 
            token=hf_token.strip() if hf_token and hf_token.strip() else None
        )
        
        # 토큰 인코딩
        tokens = tokenizer.encode(text)
        token_count = len(tokens)
        
        # 결과 반환
        result = f"✅ 토큰 수: {token_count}\n"
        result += f"모델: {model_name}\n"
        result += f"텍스트 길이: {len(text)} 글자"
        
        return result
        
    except Exception as e:
        return f"❌ 오류: {str(e)}"

def check_model(model_name, hf_token=None):
    """모델 접근 확인"""
    try:
        if not model_name:
            return "모델명을 입력해주세요."
        
        tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            token=hf_token.strip() if hf_token and hf_token.strip() else None
        )
        
        return f"✅ {model_name} 모델 접근 가능!"
        
    except Exception as e:
        return f"❌ 오류: {str(e)}"

# Gradio 인터페이스
def create_interface():
    with gr.Blocks(title="토큰 계산기") as demo:
        gr.Markdown("# 🔢 토큰 계산기")
        
        with gr.Row():
            with gr.Column():
                model_input = gr.Textbox(
                    label="모델명",
                    placeholder="예: gpt2, klue/bert-base",
                    value="gpt2"
                )
                
                token_input = gr.Textbox(
                    label="HF 토큰 (선택사항)",
                    type="password"
                )
                
                text_input = gr.Textbox(
                    label="텍스트",
                    lines=5,
                    value="안녕하세요! 테스트 텍스트입니다."
                )
                
                with gr.Row():
                    check_btn = gr.Button("모델 확인")
                    calc_btn = gr.Button("토큰 계산", variant="primary")
            
            with gr.Column():
                output = gr.Textbox(label="결과", lines=10)
        
        # 추천 모델
        gr.Markdown("### 추천 모델")
        with gr.Row():
            models = ["gpt2", "klue/bert-base", "microsoft/DialoGPT-medium"]
            for model in models:
                btn = gr.Button(model, size="sm")
                btn.click(lambda x=model: x, outputs=model_input)
        
        # 이벤트 핸들러
        check_btn.click(check_model, [model_input, token_input], output)
        calc_btn.click(count_tokens, [model_input, text_input, token_input], output)
        text_input.submit(count_tokens, [model_input, text_input, token_input], output)
    
    return demo

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