import gradio as gr
import random
import os
import pandas as pd
from huggingface_hub import InferenceClient

# 하드코딩된 언어 모델 설정
MODEL_NAME = "Cohere Command R+"
MODEL_PATH = "CohereForAI/c4ai-command-r-plus"

def create_client(model_name):
    return InferenceClient(model_name, token=os.getenv("HF_TOKEN"))

def call_api(content, system_message, max_tokens, temperature, top_p):
    client = create_client(MODEL_PATH)  # 모델을 하드코딩된 값으로 설정
    messages = [{"role": "system", "content": system_message}, {"role": "user", "content": content}]
    random_seed = random.randint(0, 1000000)
    response = client.chat_completion(messages=messages, max_tokens=max_tokens, temperature=temperature, top_p=top_p, seed=random_seed)
    return response.choices[0].message.content

def generate_text(user_message, system_message, max_tokens, temperature, top_p):
    return call_api(user_message, system_message, max_tokens, temperature, top_p)

def upload_excel(file):
    df = pd.read_excel(file.name)
    return df.head().to_string()

title = "AI 텍스트 생성기"

with gr.Blocks() as demo:
    gr.Markdown(f"# {title}")
    
    # 엑셀 업로드 기능을 가장 위로 위치
    excel_input = gr.File(label="엑셀 파일 업로드", file_types=[".xls", ".xlsx"])
    excel_output = gr.Textbox(label="엑셀 내용 미리보기", lines=10)
    
    # 사용자 메시지의 명칭을 '긍정리뷰 10개'로 설정
    user_message = gr.Textbox(label="긍정리뷰 10개", lines=5)
    
    # 입력창 1의 명칭을 '부정리뷰 10개'로 설정
    input1 = gr.Textbox(label="부정리뷰 10개", lines=5)
    
    # 시스템 메시지(프롬프트)의 명칭을 '긍정 프롬프트'로 설정
    system_message = gr.Textbox(label="긍정 프롬프트", lines=10)

    # 입력창 2의 명칭을 '부정 프롬프트'로 설정
    input2 = gr.Textbox(label="부정 프롬프트")
    
    # 출력창 1의 명칭을 '긍정리뷰분석'으로 설정
    output1 = gr.Textbox(label="긍정리뷰분석", lines=10)
    
    # 출력창 2의 명칭을 '부정리뷰분석'으로 설정
    output2 = gr.Textbox(label="부정리뷰분석", lines=10)

    with gr.Accordion("고급 설정", open=False):
        max_tokens = gr.Slider(label="Max Tokens", minimum=0, maximum=4000, value=500, step=100)
        temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.75, step=0.05)
        top_p = gr.Slider(label="Top P", minimum=0.1, maximum=1.0, value=0.95, step=0.05)
    
    # 실행 버튼 추가
    generate_btn = gr.Button("텍스트 생성하기")
    
    generate_btn.click(fn=generate_text, 
                       inputs=[user_message, system_message, max_tokens, temperature, top_p], 
                       outputs=[output1, output2])
    
    excel_input.upload(upload_excel, inputs=[excel_input], outputs=[excel_output])

demo.launch()