import os
from dotenv import load_dotenv
from langchain_huggingface import HuggingFaceEndpoint
import streamlit as st
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser

# 사용할 Hugging Face 모델 ID를 정의합니다.
model_id = "mistralai/Mistral-7B-Instruct-v0.3"

# .env 파일 로드
load_dotenv()


def get_llm_hf_inference(model_id=model_id, max_new_tokens=128, temperature=0.1):
    """
    Hugging Face 추론을 위한 언어 모델을 반환합니다.

    매개변수:
    - model_id (str): Hugging Face 모델 저장소의 ID입니다.
    - max_new_tokens (int): 생성할 수 있는 최대 새 토큰 수입니다.
    - temperature (float): 모델에서 샘플링할 때의 온도 값입니다.

    반환값:
    - llm (HuggingFaceEndpoint): Hugging Face 추론을 위한 언어 모델입니다.
    """
    # HuggingFaceEndpoint를 사용하여 언어 모델을 초기화합니다.
    llm = HuggingFaceEndpoint(
        repo_id=model_id,  # 사용할 모델 ID
        max_new_tokens=max_new_tokens,  # 생성할 최대 토큰 수
        temperature=temperature,  # 샘플링 시 온도 설정
        token=os.getenv("HF_TOKEN"),  # Hugging Face API 토큰 (환경 변수에서 가져옴)
    )
    return llm  # 초기화된 언어 모델을 반환합니다.


# Streamlit 앱 설정을 구성합니다.
st.set_page_config(page_title="HuggingFace ChatBot", page_icon="🤗")
st.title("개인 HuggingFace 챗봇")
st.markdown(
    f"*이것은 HuggingFace transformers 라이브러리를 사용하여 텍스트 입력에 대한 응답을 생성하는 간단한 챗봇입니다. {model_id} 모델을 사용합니다.*"
)

# 아바타에 대한 세션 상태를 초기화합니다.
if "avatars" not in st.session_state:
    st.session_state.avatars = {"user": None, "assistant": None}

# 사용자 텍스트 입력에 대한 세션 상태를 초기화합니다.
if "user_text" not in st.session_state:
    st.session_state.user_text = None

# 모델 매개변수에 대한 세션 상태를 초기화합니다.
if "max_response_length" not in st.session_state:
    st.session_state.max_response_length = 256

# 시스템 메시지에 대한 세션 상태를 초기화합니다.
if "system_message" not in st.session_state:
    st.session_state.system_message = "인간 사용자와 대화하는 친절한 AI"

# 시작 메시지에 대한 세션 상태를 초기화합니다.
if "starter_message" not in st.session_state:
    st.session_state.starter_message = "안녕하세요! 오늘 무엇을 도와드릴까요?"


# 설정을 위한 사이드바를 구성합니다.
with st.sidebar:
    st.header("시스템 설정")

    # AI 설정
    st.session_state.system_message = st.text_area(
        "시스템 메시지", value="당신은 인간 사용자와 대화하는 친절한 AI입니다."
    )
    st.session_state.starter_message = st.text_area(
        "첫 번째 AI 메시지", value="안녕하세요! 오늘 무엇을 도와드릴까요?"
    )

    # 모델 설정
    st.session_state.max_response_length = st.number_input("최대 응답 길이", value=128)

    # 아바타 선택
    st.markdown("*아바타 선택:*")
    col1, col2 = st.columns(2)
    with col1:
        st.session_state.avatars["assistant"] = st.selectbox(
            "AI 아바타", options=["🤗", "💬", "🤖"], index=0
        )
    with col2:
        st.session_state.avatars["user"] = st.selectbox(
            "사용자 아바타", options=["👤", "👱‍♂️", "👨🏾", "👩", "👧🏾"], index=0
        )
    # 채팅 기록 초기화 버튼
    reset_history = st.button("채팅 기록 초기화")

# 채팅 기록을 초기화하거나, 초기화 버튼이 눌렸을 경우 초기화합니다.
if "chat_history" not in st.session_state or reset_history:
    st.session_state.chat_history = [
        {"role": "assistant", "content": st.session_state.starter_message}
    ]


def get_response(
    system_message,
    chat_history,
    user_text,
    eos_token_id=["User"],
    max_new_tokens=256,
    get_llm_hf_kws={},
):
    """
    챗봇 모델로부터 응답을 생성합니다.

    매개변수:
        system_message (str): 대화의 시스템 메시지입니다.
        chat_history (list): 이전 채팅 메시지 목록입니다.
        user_text (str): 사용자의 입력 텍스트입니다.
        model_id (str, optional): 사용할 Hugging Face 모델의 ID입니다.
        eos_token_id (list, optional): 문장 종료 토큰 ID 목록입니다.
        max_new_tokens (int, optional): 생성할 수 있는 최대 새 토큰 수입니다.
        get_llm_hf_kws (dict, optional): get_llm_hf 함수에 전달할 추가 키워드 인자입니다.

    반환값:
        tuple: 생성된 응답과 업데이트된 채팅 기록을 포함하는 튜플입니다.
    """
    # 모델을 설정합니다.
    hf = get_llm_hf_inference(max_new_tokens=max_new_tokens, temperature=0.1)

    # 프롬프트 템플릿을 생성합니다.
    prompt = PromptTemplate.from_template(
        (
            "[INST] {system_message}"
            "\n현재 대화:\n{chat_history}\n\n"
            "\n사용자: {user_text}.\n [/INST]"
            "\nAI:"
        )
    )
    # 프롬프트를 연결하여 채팅 체인을 만듭니다.
    chat = prompt | hf.bind(skip_prompt=True) | StrOutputParser(output_key="content")

    # 응답을 생성합니다.
    response = chat.invoke(
        input=dict(
            system_message=system_message,
            user_text=user_text,
            chat_history=chat_history,
        )
    )
    # "AI:" 접두사를 제거합니다.
    response = response.split("AI:")[-1]

    # 채팅 기록을 업데이트합니다.
    chat_history.append({"role": "user", "content": user_text})
    chat_history.append({"role": "assistant", "content": response})
    return response, chat_history


# 채팅 인터페이스를 설정합니다.
chat_interface = st.container(border=True)
with chat_interface:
    output_container = st.container()
    st.session_state.user_text = st.chat_input(
        placeholder="여기에 텍스트를 입력하세요."
    )

# 채팅 메시지를 표시합니다.
with output_container:
    # 채팅 기록에 있는 각 메시지에 대해 반복합니다.
    for message in st.session_state.chat_history:
        # 시스템 메시지는 건너뜁니다.
        if message["role"] == "system":
            continue

        # 올바른 아바타를 사용하여 채팅 메시지를 표시합니다.
        with st.chat_message(
            message["role"], avatar=st.session_state["avatars"][message["role"]]
        ):
            st.markdown(message["content"])

    # 사용자가 새 텍스트를 입력했을 때:
    if st.session_state.user_text:
        # 사용자의 새 메시지를 즉시 표시합니다.
        with st.chat_message("user", avatar=st.session_state.avatars["user"]):
            st.markdown(st.session_state.user_text)

        # 응답을 기다리는 동안 스피너 상태 표시줄을 표시합니다.
        with st.chat_message("assistant", avatar=st.session_state.avatars["assistant"]):
            with st.spinner("생각 중..."):
                # 시스템 프롬프트, 사용자 텍스트 및 기록을 사용하여 추론 API를 호출합니다.
                response, st.session_state.chat_history = get_response(
                    system_message=st.session_state.system_message,
                    user_text=st.session_state.user_text,
                    chat_history=st.session_state.chat_history,
                    max_new_tokens=st.session_state.max_response_length,
                )
                st.markdown(response)