import gradio as gr
from huggingface_hub import InferenceClient
import os

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

def get_client(hf_token):
    """
    Cohere Command R+ 모델을 위한 InferenceClient 생성.
    hf_token은 Gradio 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_full = client.chat_completion(
            messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
        )
        assistant_message = response_full.choices[0].message.content
        return assistant_message
    except Exception as e:
        return f"오류가 발생했습니다: {str(e)}"

# 고급 설정 (코드 내에서만 정의)
SYSTEM_MESSAGE = """반드시 한글로 답변할 것.
너는 최고의 블로그 글 작성자이다.
내가 제공한 참조글을 바탕으로 블로그 글을 생성하라.
"""
MAX_TOKENS = 4000
TEMPERATURE = 0.7
TOP_P = 0.95

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

    # HuggingFace 토큰 (secrets 사용)
    hf_token = os.getenv("HF_TOKEN")

    # 입력 필드
    with gr.Row():
        tone = gr.Radio(
            choices=["친근하게", "일반적인", "전문적인"],
            label="말투바꾸기",
            value="일반적인"
        )
    ref1 = gr.Textbox(label="참조글 1", lines=3)
    ref2 = gr.Textbox(label="참조글 2", lines=3)
    ref3 = gr.Textbox(label="참조글 3", lines=3)

    # 결과 출력
    answer_output = gr.Textbox(label="생성된 블로그 글", lines=10, interactive=False)

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

    def generate_blog(tone, ref1, ref2, ref3):
        # HuggingFace 토큰은 함수 내부에서 직접 사용
        hf_token_value = os.getenv("HF_TOKEN")
        if not hf_token_value:
            return "HuggingFace 토큰이 설정되지 않았습니다."

        # 프롬프트 구성
        prompt = f"말투: {tone}\n참조글 1: {ref1}\n참조글 2: {ref2}\n참조글 3: {ref3}"
        return respond_cohere_qna(
            question=prompt,
            system_message=SYSTEM_MESSAGE,
            max_tokens=MAX_TOKENS,
            temperature=TEMPERATURE,
            top_p=TOP_P,
            hf_token=hf_token_value
        )

    submit_button.click(
        fn=generate_blog,
        inputs=[tone, ref1, ref2, ref3],  # hf_token은 제외
        outputs=answer_output
    )

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