import gradio as gr
from huggingface_hub import InferenceClient
import os
from typing import Optional

# Cohere Command R+ 모델 ID
COHERE_MODEL = "CohereForAI/c4ai-command-r-plus-08-2024"

def get_client(hf_token):
    """
    HuggingFace InferenceClient 생성.
    hf_token을 secrets에서 가져옴.
    """
    if not hf_token:
        raise ValueError("HuggingFace API 토큰이 필요합니다.")
    return InferenceClient(COHERE_MODEL, token=hf_token)

def respond_cohere_qna(
    question: str,
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    hf_token: str
):
    """
    Cohere Command R+ 모델을 이용해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
    """
    try:
        client = get_client(hf_token)
    except ValueError as e:
        return f"오류: {str(e)}"

    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": question}
    ]

    try:
        response = client.chat_completion(
            messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            stream=False,
        )
        assistant_message = response.choices[0].message.content
        return assistant_message
    except Exception as e:
        return f"오류가 발생했습니다: {str(e)}"

# Gradio UI 구성
with gr.Blocks() as demo:
    gr.Markdown("# 블로그 생성기")

    # 입력1: 말투 바꾸기 (라디오 버튼)
    tone_radio = gr.Radio(
        choices=["친근하게", "일반적인", "전문적인"],
        label="말투 바꾸기",
        value="일반적인"
    )

    # 입력2: 참조글 1
    ref1 = gr.Textbox(label="참조글 1", lines=3)

    # 입력3: 참조글 2
    ref2 = gr.Textbox(label="참조글 2", lines=3)

    # 입력4: 참조글 3
    ref3 = gr.Textbox(label="참조글 3", lines=3)

    # 입력5: 내 글 내용
    my_content = gr.Textbox(label="내 글 내용", lines=3, placeholder="내가 쓰고 싶은 주제를 입력하세요.")

    # 결과 출력
    answer_output = gr.Textbox(label="결과", lines=10, interactive=False)

    # 고급 설정 (코드 내부에서만 사용)
    system_message = gr.Textbox(
        value="""반드시 한글로 답변할 것.
너는 최고의 블로그 글 생성기이다.
주어진 참조글과 내 글 내용을 바탕으로 블로그 글을 생성하라.
말투는 선택된 옵션에 맞게 조정하라.
""",
        visible=False  # UI에서 숨김
    )
    max_tokens = gr.Slider(minimum=100, maximum=10000, value=4000, step=100, label="Max Tokens", visible=False)
    temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature", visible=False)
    top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P", visible=False)

    # 전송 버튼
    submit_button = gr.Button("생성")

    def merge_and_call_cohere(tone, ref1, ref2, ref3, my_content, sys_msg, mt, temp, top_p_):
        # 참조글과 내 글 내용을 합쳐서 질문 구성
        question = f"말투: {tone}\n참조글 1: {ref1}\n참조글 2: {ref2}\n참조글 3: {ref3}\n내 글 내용: {my_content}"
        hf_token = os.environ.get("HF_TOKEN")
        return respond_cohere_qna(
            question=question,
            system_message=sys_msg,
            max_tokens=mt,
            temperature=temp,
            top_p=top_p_,
            hf_token=hf_token
        )

    submit_button.click(
        fn=merge_and_call_cohere,
        inputs=[
            tone_radio, ref1, ref2, ref3, my_content,
            system_message,
            max_tokens,
            temperature,
            top_p
        ],
        outputs=answer_output
    )

# 메인 실행부
if __name__ == "__main__":
    demo.launch()