import gradio as gr
from huggingface_hub import InferenceClient
import openai
import anthropic
import os
from typing import Optional
import transformers
import torch

#############################
# [기본코드] - 수정/삭제 불가
#############################

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

def get_client(model_name):
    """
    모델 이름에 맞춰 InferenceClient 생성.
    토큰은 환경 변수에서 가져옴.
    """
    hf_token = os.getenv("HF_TOKEN")
    if not hf_token:
        raise ValueError("HuggingFace API 토큰이 필요합니다.")

    if model_name == "Cohere Command R+":
        model_id = COHERE_MODEL
    else:
        raise ValueError("유효하지 않은 모델 이름입니다.")
    return InferenceClient(model_id, token=hf_token)

def respond_cohere_qna(
    question: str,
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float
):
    """
    Cohere Command R+ 모델을 이용해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
    """
    model_name = "Cohere Command R+"
    try:
        client = get_client(model_name)
    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)}"

def respond_chatgpt_qna(
    question: str,
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float
):
    """
    ChatGPT(OpenAI) 모델을 이용해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
    """
    openai_token = os.getenv("OPENAI_TOKEN")
    if not openai_token:
        return "OpenAI API 토큰이 필요합니다."

    openai.api_key = openai_token

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

    try:
        response = openai.ChatCompletion.create(
            model="gpt-4o-mini",  # 필요한 경우 변경
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
        )
        assistant_message = response.choices[0].message['content']
        return assistant_message
    except Exception as e:
        return f"오류가 발생했습니다: {str(e)}"

def respond_deepseek_qna(
    question: str,
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    model_name: str  # 모델 이름 추가
):
    """
    DeepSeek 모델을 이용해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
    """
    deepseek_token = os.getenv("DEEPSEEK_TOKEN")
    if not deepseek_token:
        return "DeepSeek API 토큰이 필요합니다."

    openai.api_key = deepseek_token
    openai.api_base = "https://api.deepseek.com/v1"

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

    try:
        response = openai.ChatCompletion.create(
            model=model_name,  # 선택된 모델 사용
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
        )
        assistant_message = response.choices[0].message['content']
        return assistant_message
    except Exception as e:
        return f"오류가 발생했습니다: {str(e)}"

def respond_claude_qna(
    question: str,
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float
) -> str:
    """
    Claude API를 사용한 개선된 응답 생성 함수
    """
    claude_api_key = os.getenv("CLAUDE_TOKEN")
    if not claude_api_key:
        return "Claude API 토큰이 필요합니다."

    try:
        client = anthropic.Anthropic(api_key=claude_api_key)
        
        # 메시지 생성
        message = client.messages.create(
            model="claude-3-haiku-20240307",
            max_tokens=max_tokens,
            temperature=temperature,
            system=system_message,
            messages=[
                {
                    "role": "user",
                    "content": question
                }
            ]
        )
        
        return message.content[0].text
            
    except anthropic.APIError as ae:
        return f"Claude API 오류: {str(ae)}"
    except anthropic.RateLimitError:
        return "요청 한도를 초과했습니다. 잠시 후 다시 시도해주세요."
    except Exception as e:
        return f"예상치 못한 오류가 발생했습니다: {str(e)}"

#############################
# [추가코드] - Llama-3.3-70B-Instruct / Llama-3.2-3B-Instruct 적용 (transformers.pipeline 방식)
#############################

def get_llama_client(model_choice: str):
    """
    선택된 Llama 모델에 맞춰 transformers의 text-generation 파이프라인을 생성.
    """
    if model_choice == "Llama-3.3-70B-Instruct":
        model_id = "meta-llama/Llama-3.3-70B-Instruct"
    elif model_choice == "Llama-3.2-3B-Instruct":
        model_id = "meta-llama/Llama-3.2-3B-Instruct"
    else:
        raise ValueError("유효하지 않은 모델 선택입니다.")
    
    pipeline_llama = transformers.pipeline(
        "text-generation",
        model=model_id,
        model_kwargs={"torch_dtype": torch.bfloat16},
        device_map="auto",
    )
    return pipeline_llama

def respond_llama_qna(
    question: str,
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    model_choice: str
):
    """
    선택된 Llama 모델을 이용해 한 번의 질문(question)에 대한 답변을 transformers 파이프라인으로 반환하는 함수.
    system_message와 question을 하나의 프롬프트로 결합하여 생성합니다.
    """
    try:
        pipeline_llama = get_llama_client(model_choice)
    except ValueError as e:
        return f"오류: {str(e)}"
    
    # system_message와 question을 연결하여 프롬프트 생성
    prompt = system_message.strip() + "\n" + question.strip()
    
    try:
        outputs = pipeline_llama(
            prompt,
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
        )
        # 생성된 텍스트를 추출 (전체 프롬프트 이후의 텍스트만 반환할 수도 있음)
        generated_text = outputs[0]["generated_text"]
        return generated_text
    except Exception as e:
        return f"오류가 발생했습니다: {str(e)}"

#############################
# [기본코드] UI 부분 - 수정/삭제 불가
#############################

with gr.Blocks() as demo:
    gr.Markdown("# LLM 플레이그라운드")

    #################
    # Cohere Command R+ 탭
    #################
    with gr.Tab("Cohere Command R+"):
        cohere_input1 = gr.Textbox(label="입력1", lines=1)
        cohere_input2 = gr.Textbox(label="입력2", lines=1)
        cohere_input3 = gr.Textbox(label="입력3", lines=1)
        cohere_input4 = gr.Textbox(label="입력4", lines=1)
        cohere_input5 = gr.Textbox(label="입력5", lines=1)

        cohere_answer_output = gr.Textbox(label="결과", lines=5, interactive=False)

        with gr.Accordion("고급 설정 (Cohere)", open=False):
            cohere_system_message = gr.Textbox(
                value="""반드시 한글로 답변할 것.
너는 최고의 비서이다.
내가 요구하는것들을 최대한 자세하고 정확하게 답변하라.
""",
                label="System Message",
                lines=3
            )
            cohere_max_tokens = gr.Slider(minimum=100, maximum=10000, value=4000, step=100, label="Max Tokens")
            cohere_temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
            cohere_top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")

        cohere_submit_button = gr.Button("전송")

        def merge_and_call_cohere(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_):
            question = " ".join([i1, i2, i3, i4, i5])
            return respond_cohere_qna(
                question=question,
                system_message=sys_msg,
                max_tokens=mt,
                temperature=temp,
                top_p=top_p_
            )

        cohere_submit_button.click(
            fn=merge_and_call_cohere,
            inputs=[
                cohere_input1, cohere_input2, cohere_input3, cohere_input4, cohere_input5,
                cohere_system_message,
                cohere_max_tokens,
                cohere_temperature,
                cohere_top_p
            ],
            outputs=cohere_answer_output
        )

    #################
    # ChatGPT 탭
    #################
    with gr.Tab("gpt-4o-mini"):
        chatgpt_input1 = gr.Textbox(label="입력1", lines=1)
        chatgpt_input2 = gr.Textbox(label="입력2", lines=1)
        chatgpt_input3 = gr.Textbox(label="입력3", lines=1)
        chatgpt_input4 = gr.Textbox(label="입력4", lines=1)
        chatgpt_input5 = gr.Textbox(label="입력5", lines=1)

        chatgpt_answer_output = gr.Textbox(label="결과", lines=5, interactive=False)

        with gr.Accordion("고급 설정 (ChatGPT)", open=False):
            chatgpt_system_message = gr.Textbox(
                value="""반드시 한글로 답변할 것.
너는 ChatGPT, OpenAI에서 개발한 언어 모델이다.
내가 요구하는 것을 최대한 자세하고 정확하게 답변하라.
""",
                label="System Message",
                lines=3
            )
            chatgpt_max_tokens = gr.Slider(minimum=100, maximum=4000, value=2000, step=100, label="Max Tokens")
            chatgpt_temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="Temperature")
            chatgpt_top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")

        chatgpt_submit_button = gr.Button("전송")

        def merge_and_call_chatgpt(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_):
            question = " ".join([i1, i2, i3, i4, i5])
            return respond_chatgpt_qna(
                question=question,
                system_message=sys_msg,
                max_tokens=mt,
                temperature=temp,
                top_p=top_p_
            )

        chatgpt_submit_button.click(
            fn=merge_and_call_chatgpt,
            inputs=[
                chatgpt_input1, chatgpt_input2, chatgpt_input3, chatgpt_input4, chatgpt_input5,
                chatgpt_system_message,
                chatgpt_max_tokens,
                chatgpt_temperature,
                chatgpt_top_p
            ],
            outputs=chatgpt_answer_output
        )

    #################
    # Claude 탭
    #################
    with gr.Tab("claude-3-haiku"):
        claude_input1 = gr.Textbox(label="입력1", lines=1)
        claude_input2 = gr.Textbox(label="입력2", lines=1)
        claude_input3 = gr.Textbox(label="입력3", lines=1)
        claude_input4 = gr.Textbox(label="입력4", lines=1)
        claude_input5 = gr.Textbox(label="입력5", lines=1)

        claude_answer_output = gr.Textbox(label="결과", interactive=False, lines=5)

        with gr.Accordion("고급 설정 (Claude)", open=False):
            claude_system_message = gr.Textbox(
                label="System Message",
                value="""반드시 한글로 답변할 것.
너는 Anthropic에서 개발한 클로드이다.
최대한 정확하고 친절하게 답변하라.""",
                lines=3
            )
            claude_max_tokens = gr.Slider(
                minimum=100,
                maximum=4000,
                value=2000,
                step=100,
                label="Max Tokens"
            )
            claude_temperature = gr.Slider(
                minimum=0.1,
                maximum=2.0,
                value=0.7,
                step=0.05,
                label="Temperature"
            )
            claude_top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.95,
                step=0.05,
                label="Top-p"
            )

        claude_submit_button = gr.Button("전송")

        def merge_and_call_claude(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_):
            question = " ".join([i1, i2, i3, i4, i5])
            return respond_claude_qna(
                question=question,
                system_message=sys_msg,
                max_tokens=mt,
                temperature=temp,
                top_p=top_p_
            )

        claude_submit_button.click(
            fn=merge_and_call_claude,
            inputs=[
                claude_input1, claude_input2, claude_input3, claude_input4, claude_input5,
                claude_system_message,
                claude_max_tokens,
                claude_temperature,
                claude_top_p
            ],
            outputs=claude_answer_output
        )

    #################
    # DeepSeek 탭
    #################
    with gr.Tab("DeepSeek-V3"):
        # 라디오 버튼 추가
        deepseek_model_radio = gr.Radio(
            choices=["V3 (deepseek-chat)", "R1 (deepseek-reasoner)"],  # 선택지
            label="모델 선택",  # 라벨
            value="V3 (deepseek-chat)"  # 기본값
        )

        deepseek_input1 = gr.Textbox(label="입력1", lines=1)
        deepseek_input2 = gr.Textbox(label="입력2", lines=1)
        deepseek_input3 = gr.Textbox(label="입력3", lines=1)
        deepseek_input4 = gr.Textbox(label="입력4", lines=1)
        deepseek_input5 = gr.Textbox(label="입력5", lines=1)

        deepseek_answer_output = gr.Textbox(label="결과", lines=5, interactive=False)

        with gr.Accordion("고급 설정 (DeepSeek)", open=False):
            deepseek_system_message = gr.Textbox(
                value="""반드시 한글로 답변할 것.
너는 DeepSeek-V3, 최고의 언어 모델이다.
내가 요구하는 것을 최대한 자세하고 정확하게 답변하라.
""",
                label="System Message",
                lines=3
            )
            deepseek_max_tokens = gr.Slider(minimum=100, maximum=4000, value=2000, step=100, label="Max Tokens")
            deepseek_temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="Temperature")
            deepseek_top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")

        deepseek_submit_button = gr.Button("전송")

        def merge_and_call_deepseek(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_, model_radio):
            # 라디오 버튼에서 선택된 모델 이름 추출
            if model_radio == "V3 (deepseek-chat)":
                model_name = "deepseek-chat"
            else:
                model_name = "deepseek-reasoner"

            question = " ".join([i1, i2, i3, i4, i5])
            return respond_deepseek_qna(
                question=question,
                system_message=sys_msg,
                max_tokens=mt,
                temperature=temp,
                top_p=top_p_,
                model_name=model_name  # 선택된 모델 이름 전달
            )

        deepseek_submit_button.click(
            fn=merge_and_call_deepseek,
            inputs=[
                deepseek_input1, deepseek_input2, deepseek_input3, deepseek_input4, deepseek_input5,
                deepseek_system_message,
                deepseek_max_tokens,
                deepseek_temperature,
                deepseek_top_p,
                deepseek_model_radio  # 라디오 버튼 입력 추가
            ],
            outputs=deepseek_answer_output
        )

    #################
    # Llama 탭 (추가)
    #################
    with gr.Tab("Llama"):
        # 라디오 버튼 추가: Llama-3.3-70B-Instruct (기본) / Llama-3.2-3B-Instruct
        llama_model_radio = gr.Radio(
            choices=["Llama-3.3-70B-Instruct", "Llama-3.2-3B-Instruct"],
            label="모델 선택",
            value="Llama-3.3-70B-Instruct"
        )

        llama_input1 = gr.Textbox(label="입력1", lines=1)
        llama_input2 = gr.Textbox(label="입력2", lines=1)
        llama_input3 = gr.Textbox(label="입력3", lines=1)
        llama_input4 = gr.Textbox(label="입력4", lines=1)
        llama_input5 = gr.Textbox(label="입력5", lines=1)

        llama_answer_output = gr.Textbox(label="결과", lines=5, interactive=False)

        with gr.Accordion("고급 설정 (Llama)", open=False):
            llama_system_message = gr.Textbox(
                value="""반드시 한글로 답변할 것.
너는 최고의 비서이다.
내가 요구하는 것을 최대한 자세하고 정확하게 답변하라.
""",
                label="System Message",
                lines=3
            )
            llama_max_tokens = gr.Slider(minimum=100, maximum=10000, value=4000, step=100, label="Max Tokens")
            llama_temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
            llama_top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")

        llama_submit_button = gr.Button("전송")

        def merge_and_call_llama(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_, model_choice):
            question = " ".join([i1, i2, i3, i4, i5])
            return respond_llama_qna(
                question=question,
                system_message=sys_msg,
                max_tokens=mt,
                temperature=temp,
                top_p=top_p_,
                model_choice=model_choice
            )

        llama_submit_button.click(
            fn=merge_and_call_llama,
            inputs=[
                llama_input1, llama_input2, llama_input3, llama_input4, llama_input5,
                llama_system_message,
                llama_max_tokens,
                llama_temperature,
                llama_top_p,
                llama_model_radio  # 라디오 버튼 입력 추가
            ],
            outputs=llama_answer_output
        )

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