# -*- coding: utf-8 -*-

import gradio as gr
from huggingface_hub import InferenceClient
import os
import requests

# 추론 API 클라이언트 설정
hf_client = InferenceClient("CohereForAI/c4ai-command-r-plus-08-2024", token=os.getenv("HF_TOKEN"))

def respond(
    message,
    history: list[tuple[str, str]],
    system_message="",  
    max_tokens=7860,  
    temperature=0.8, 
    top_p=0.9,  
):
    system_prefix = """
    [시스템 프롬프트 내용...]
    """

    messages = [{"role": "system", "content": f"{system_prefix} {system_message}"}]
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})
    messages.append({"role": "user", "content": message})

    response = ""
    try:
        for message in hf_client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        ):
            token = message.choices[0].delta.content
            if token is not None:
                response += token.strip("")
            yield response
    except Exception as e:
        yield f"Error: {str(e)}"

# Gradio 인터페이스 설정
interface = gr.ChatInterface(
    respond,    
    additional_inputs=[
        gr.Textbox(label="System Message", value="Write(output) in 한국어."),
        gr.Slider(minimum=1, maximum=8000, value=7000, label="Max Tokens"),
        gr.Slider(minimum=0, maximum=1, value=0.7, label="Temperature"),
        gr.Slider(minimum=0, maximum=1, value=0.9, label="Top P"),
    ],
    examples=[
        ["판타지 소설의 흥미로운 소재 10가지를 제시하라"],
        ["계속 이어서 작성하라"],
        ["Translate into English"],
        ["마법 시스템에 대해 더 자세히 설명하라"],
        ["전투 장면을 더 극적으로 묘사하라"],
        ["새로운 판타지 종족을 추가하라"],
        ["고대 예언에 대해 더 자세히 설명하라"],
        ["주인공의 내면 묘사를 추가하라"],
    ],
    title="Fantasy Novel AI Generation",  
    cache_examples=False,
    theme="Yntec/HaleyCH_Theme_Orange"
)

# 애플리케이션 실행
if __name__ == "__main__":
    interface.launch(
        server_name="0.0.0.0",  # 모든 IP에서 접근 가능
        server_port=7860,        # 포트 지정
        share=True              # 공유 링크 생성
    )