import streamlit as st
from together import Together
import os
from typing import Iterator
from PIL import Image
import base64
from PyPDF2 import PdfReader

API_KEY = os.getenv("TOGETHER_API_KEY")
if not API_KEY:
    raise ValueError("API key is missing! Make sure TOGETHER_API_KEY is set in the Secrets.")

@st.cache_resource
def get_client():
    return Together(api_key=API_KEY)

def process_file(file) -> str:
    if file is None:
        return ""
    
    try:
        if file.type == "application/pdf":
            text = ""
            pdf_reader = PdfReader(file)
            for page in pdf_reader.pages:
                text += page.extract_text() + "\n"
            return text
        elif file.type.startswith("image/"):
            return base64.b64encode(file.getvalue()).decode("utf-8")
        else:
            return file.getvalue().decode('utf-8')
    except Exception as e:
        st.error(f"파일 처리 중 오류 발생: {str(e)}")
        return ""

def generate_response(
    message: str,
    history: list[tuple[str, str]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    files=None
) -> Iterator[str]:
    client = get_client()
    
    try:
        # 메시지 배열 초기화
        messages = []
        
        # 시스템 메시지가 있는 경우에만 추가
        if system_message.strip():
            messages.append({
                "role": "system",
                "content": system_message
            })
        
        # 대화 히스토리 추가
        for user_msg, assistant_msg in history:
            messages.append({
                "role": "user",
                "content": user_msg
            })
            messages.append({
                "role": "assistant",
                "content": assistant_msg
            })
        
        # 현재 메시지와 파일 내용 준비
        current_content = message
        if files:
            file_contents = []
            for file in files:
                content = process_file(file)
                if content:
                    file_contents.append(f"파일 내용:\n{content}")
            if file_contents:
                current_content = current_content + "\n\n" + "\n\n".join(file_contents)
        
        # 현재 메시지 추가
        messages.append({
            "role": "user",
            "content": current_content
        })
        
        # API 요청 설정
        request_params = {
            "model": "deepseek-ai/DeepSeek-R1",
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "stream": True
        }
        
        # API 호출
        try:
            stream = client.chat.completions.create(**request_params)
            
            for chunk in stream:
                if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
                    yield chunk.choices[0].delta.content
                    
        except Exception as e:
            if "rate limit" in str(e).lower():
                yield "API 호출 한도에 도달했습니다. 잠시 후 다시 시도해주세요."
            else:
                error_message = str(e)
                # Together.ai의 오류 응답 분석
                if "Input validation error" in error_message:
                    yield "입력 형식이 올바르지 않습니다. 시스템 관리자에게 문의해주세요."
                else:
                    yield f"API 호출 중 오류가 발생했습니다: {error_message}"
                
    except Exception as e:
        yield f"오류가 발생했습니다: {str(e)}"

def main():
    st.set_page_config(page_title="DeepSeek 채팅", page_icon="💭", layout="wide")
    
    if "messages" not in st.session_state:
        st.session_state.messages = []

    st.title("DeepSeek 채팅")
    st.markdown("DeepSeek AI 모델과 대화하세요. 필요한 경우 파일을 업로드할 수 있습니다.")

    with st.sidebar:
        st.header("설정")
        system_message = st.text_area(
            "시스템 메시지",
            value="당신은 깊이 있게 생각하는 AI입니다. 문제를 깊이 고려하고 체계적인 추론 과정을 통해 올바른 해결책을 도출하세요. 반드시 한글로 답변하세요.",
            height=100
        )
        max_tokens = st.slider("최대 토큰 수", 1, 4096, 2048)  # 토큰 제한 조정
        temperature = st.slider("온도", 0.0, 2.0, 0.7, 0.1)    # 온도 범위 조정
        top_p = st.slider("Top-p", 0.0, 1.0, 0.7, 0.1)        # top_p 범위 조정
        uploaded_file = st.file_uploader(
            "파일 업로드 (선택사항)",
            type=['txt', 'py', 'md', 'pdf', 'png', 'jpg', 'jpeg'],
            accept_multiple_files=True
        )

    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    if prompt := st.chat_input("무엇을 알고 싶으신가요?"):
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        with st.chat_message("assistant"):
            response_placeholder = st.empty()
            full_response = ""
            
            history = [(msg["content"], next_msg["content"]) 
                      for msg, next_msg in zip(st.session_state.messages[::2], st.session_state.messages[1::2])]
            
            for response_chunk in generate_response(
                prompt,
                history,
                system_message,
                max_tokens,
                temperature,
                top_p,
                uploaded_file
            ):
                full_response += response_chunk
                response_placeholder.markdown(full_response + "▌")
            
            response_placeholder.markdown(full_response)
        
        st.session_state.messages.append({"role": "assistant", "content": full_response})

if __name__ == "__main__":
    main()