import streamlit as st
import requests

# -------------------------------
# 1. 문서 파싱 함수 (HTML 추출)
# -------------------------------
def parse_document_to_html(file, api_key):
    url = "https://api.upstage.ai/v1/document-ai/document-parse"
    headers = {'Authorization': f'Bearer {api_key}'}
    files = {"document": file}
    data = {
        "base64_encoding": "['table']",
        "model": "document-parse"
    }

    response = requests.post(url, headers=headers, files=files, data=data)
    if response.status_code != 200:
        return None, f"문서 파싱 실패 (status code: {response.status_code})"
    
    html_text = response.json().get("content", {}).get("html", "")
    return html_text, None

# -------------------------------
# 2. Solar-Pro: 질문 응답
# -------------------------------
def chat_with_solar_document(html_text, user_question, history, api_key):
    url = "https://api.upstage.ai/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    system_prompt = f"""
    다음은 HTML 형식으로 추출된 문서입니다.
    HTML 안의 내용만 근거로 삼아 학생의 질문에 정확하고 쉽게 답해주세요.

    문서:
    {html_text}
    """

    messages = [{"role": "system", "content": system_prompt}]
    for user, bot in history:
        messages.append({"role": "user", "content": user})
        messages.append({"role": "assistant", "content": bot})
    messages.append({"role": "user", "content": user_question})

    payload = {
        "model": "solar-pro",
        "messages": messages,
        "temperature": 0,
        "max_tokens": 2048
    }

    try:
        response = requests.post(url, headers=headers, json=payload)
        if response.status_code != 200:
            return None, history, f"상태 코드: {response.status_code}"
        bot_reply = response.json()["choices"][0]["message"]["content"]
        history.append((user_question, bot_reply))
        return bot_reply, history, None
    except Exception as e:
        return None, history, f"예외 발생: {str(e)}"

# -------------------------------
# 3. 추천 키워드 생성
# -------------------------------
def get_recommended_keywords(html_text, user_question, api_key):
    url = "https://api.upstage.ai/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    prompt = f"""
    다음은 HTML로 추출된 문제 내용입니다.
    문제와 질문을 바탕으로 학생에게 도움이 될 만한 공부 키워드 3가지를 제안해주세요. 키워드만 콤마로 구분해서 출력해주세요.

    문서: {html_text}

    학생 질문: {user_question}
    """

    payload = {
        "model": "solar-pro",
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.5,
        "max_tokens": 100
    }

    try:
        res = requests.post(url, headers=headers, json=payload)
        if res.status_code != 200:
            return []
        raw_text = res.json()["choices"][0]["message"]["content"]
        keywords = [kw.strip() for kw in raw_text.split(",") if kw.strip()]
        return keywords[:3]
    except:
        return []

# -------------------------------
# 4. 키워드별 공부 팁 생성
# -------------------------------
def get_study_tip_for_keyword(keyword, api_key):
    url = "https://api.upstage.ai/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    prompt = f"""
    학생이 '{keyword}'을(를) 공부하고 싶어합니다.
    개념 이해, 공부 방법, 추천 순서 등을 포함해 친절하게 학습 팁을 제공해주세요.
    """

    payload = {
        "model": "solar-pro",
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.7,
        "max_tokens": 300
    }

    try:
        res = requests.post(url, headers=headers, json=payload)
        return res.json()["choices"][0]["message"]["content"]
    except:
        return "⚠️ 학습 추천 생성에 실패했어요."

# -------------------------------
# 5. Streamlit UI
# -------------------------------
def main():
    st.set_page_config(page_title="Solar 공부 멘토", layout="wide")
    st.title("👩‍🏫 Solar 공부 멘토")

    api_key = st.text_input("🔐 Upstage API Key", type="password")
    uploaded_file = st.file_uploader("📎 문제 이미지 업로드", type=["pdf", "png", "jpg", "jpeg"])

    if "html_text" not in st.session_state:
        st.session_state.html_text = ""
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []
    if "keywords" not in st.session_state:
        st.session_state.keywords = []
    if "study_tips" not in st.session_state:
        st.session_state.study_tips = {}

    if uploaded_file and api_key and st.button("📄 문서 분석 시작"):
        with st.spinner("문서를 파싱하고 있어요..."):
            html_text, err = parse_document_to_html(uploaded_file, api_key)
            if err:
                st.error(err)
                return
            st.session_state.html_text = html_text
            st.success("문서 파싱 성공 ✅")

    # ✨ 질의응답: 이미지와 분리해서 별도 영역에 위치
    if st.session_state.html_text:
        left_col, right_col = st.columns([1, 2])

        with left_col:
            st.image(uploaded_file, caption="🖼️ 업로드한 문제 이미지", use_container_width=True)

        with right_col:
            st.subheader("❓ 궁금한 내용을 질문해보세요")
            user_question = st.text_input("질문을 입력하고 아래 버튼을 눌러보세요", key="question_input")

            if st.button("💬 Solar에게 물어보기") and user_question:
                with st.spinner("Solar 선생님이 답변 중입니다..."):
                    answer, history, err = chat_with_solar_document(
                        st.session_state.html_text,
                        user_question,
                        st.session_state.chat_history,
                        api_key
                    )
                    if err:
                        st.error(f"❌ Solar 응답 오류: {err}")
                    else:
                        st.session_state.chat_history = history
                        st.chat_message("user").write(user_question)
                        st.chat_message("assistant").write(answer)

                        # 키워드 추천 생성
                        st.session_state.keywords = get_recommended_keywords(
                            st.session_state.html_text,
                            user_question,
                            api_key
                        )

    # ✨ 이미지 아래에 추천 키워드 표시
    if st.session_state.keywords:
        st.markdown("### 🔎 추천 공부 키워드")
        for keyword in st.session_state.keywords:
            if st.button(f"📘 {keyword} 공부 추천받기", key=f"btn_{keyword}"):
                tip = get_study_tip_for_keyword(keyword, api_key)
                st.session_state.study_tips[keyword] = tip

        for keyword, tip in st.session_state.study_tips.items():
            with st.expander(f"📘 {keyword} 공부법 보기"):
                st.write(tip)


if __name__ == "__main__":
    main()