import streamlit as st
import logging
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# 로깅 설정
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 페이지 설정
st.set_page_config(
    page_title="DeepSeek Chatbot - ruslanmv.com",
    page_icon="🤖",
    layout="centered"
)

# 세션 상태에 채팅 기록 초기화
if "messages" not in st.session_state:
    st.session_state.messages = []

# 사이드바 설정
with st.sidebar:
    st.header("Model Configuration")
    st.markdown("모델을 로컬에서 직접 로드합니다.")
    
    # 모델 선택 드롭다운 (필요 시 다른 모델 추가)
    model_options = [
        "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
    ]
    selected_model = st.selectbox("Select Model", model_options, index=0)
    
    # 시스템 메시지 설정
    system_message = st.text_area(
        "System Message",
        value="You are a friendly chatbot created by ruslanmv.com. Provide clear, accurate, and brief answers. Keep responses polite, engaging, and to the point. If unsure, politely suggest alternatives.",
        height=100
    )
    
    # 생성 파라미터 설정
    max_tokens = st.slider("Max Tokens", 10, 4000, 1000)
    temperature = st.slider("Temperature", 0.1, 4.0, 0.3)
    top_p = st.slider("Top-p", 0.1, 1.0, 0.6)

# 모델과 토크나이저를 로드하는 함수 (캐싱하여 한 번만 로드)
@st.cache_resource
def load_model_and_tokenizer(model_name: str):
    logger.info(f"Loading model and tokenizer for {model_name} ...")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    # device_map="auto" 옵션은 사용 가능한 GPU/CPU에 맞게 모델을 자동으로 할당합니다.
    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
    return tokenizer, model

tokenizer, model = load_model_and_tokenizer(selected_model)

# 채팅 인터페이스
st.title("🤖 DeepSeek Chatbot")
st.caption("Powered by local model - Configure in sidebar")

# 기존 채팅 기록 표시
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# 사용자 입력 처리
if prompt := st.chat_input("Type your message..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    with st.chat_message("user"):
        st.markdown(prompt)
    
    try:
        with st.spinner("Generating response..."):
            # 시스템 메시지와 사용자 입력을 결합하여 전체 프롬프트 구성
            full_prompt = f"{system_message}\n\nUser: {prompt}\nAssistant:"
            inputs = tokenizer.encode(full_prompt, return_tensors="pt").to(model.device)
            
            # 모델을 사용하여 응답 생성
            output_tokens = model.generate(
                inputs,
                max_new_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
            )
            
            # 생성된 토큰 디코딩
            output_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
            # "Assistant:" 이후의 텍스트만 추출 (없으면 전체 텍스트 사용)
            if "Assistant:" in output_text:
                assistant_response = output_text.split("Assistant:")[-1].strip()
            else:
                assistant_response = output_text.strip()
            
            logger.info(f"Generated response: {assistant_response}")
            
            # 생성된 응답을 채팅에 출력
            with st.chat_message("assistant"):
                st.markdown(assistant_response)
            
            st.session_state.messages.append({"role": "assistant", "content": assistant_response})
    
    except Exception as e:
        logger.error(f"Application Error: {str(e)}", exc_info=True)
        st.error(f"Application Error: {str(e)}")