import os
import time
import streamlit as st
import google.generativeai as genai
from streamlit_extras.colored_header import colored_header
from streamlit_extras.add_vertical_space import add_vertical_space

# Google Gemini API Key 설정
genai.configure(api_key=os.environ["GEMINI_API_KEY"])

# 모델 설정
generation_config = {
    "temperature": 1,
    "top_p": 0.95,
    "top_k": 40,
    "max_output_tokens": 8192,
    "response_mime_type": "text/plain",
}

model = genai.GenerativeModel(
    model_name="gemini-2.0-flash",
    generation_config=generation_config,
)

def generate_research_plan(question):
    """
    질문이나 탐구주제를 입력받아 통합탐구 계획을 생성하는 함수
    """

    # 전체 프롬프트 생성
    prompt = f"""학생들이 질문이나 탐구주제를 입력하면 과학자로서 통합탐구의 절차에 따라 실험을 구체적, 단계적, 체계적으로 설계하고 예상되는 결과를 설명해줘.

질문 및 탐구 주제: {question}

통합탐구 과정:

1. 문제 인식:

2. 가설 설정:

3. 변인 통제:

4. 자료 해석 방법 및 절차:

5. 결론 도출 방법 및 절차:

6. 예상되는 일반화:

7. 탐구를 위한 팁:"""

    # 스트리밍 방식으로 응답 생성
    response = model.generate_content([prompt], stream=True)
    collected_text = ""
    for token in response:
        chunk = token.text
        collected_text += chunk
        yield collected_text  # 실시간 출력을 위해 yield 사용
        time.sleep(0.05)  # 출력 속도 조절

# Streamlit Interface
colored_header(
    label="통합탐구 계획 생성기",
    description="과학 탐구 주제에 대한 자세한 계획을 세워보세요!",
    color_name="blue-70",
)

add_vertical_space(1)

question = st.text_area("질문이나 탐구주제를 입력하세요", height=100)

generate_button = st.button("통합탐구 계획 생성")

# 출력 영역 정의
output_area = st.empty()

if generate_button and question:
    output_text = ""
    output_area.markdown(output_text)
    try:
        for partial_output in generate_research_plan(question):
            output_text = partial_output
            output_area.markdown(output_text)
    except Exception as e:
        st.error(f"오류가 발생했습니다: {str(e)}")

    # 복사 버튼 추가
    if output_text:
        if st.button("출력 내용 복사"):
            st.write(output_text)
            st.success("출력 내용이 복사되었습니다!")