import os
import time
import google.generativeai as genai
import gradio as gr

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",
}
safety_settings = [
    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
]

# 전역 변수로 chat_session 초기화
chat_session = None


def start_new_chat_session():
    global chat_session
    chat_session = (
        genai.GenerativeModel(
            model_name="gemini-2.0-flash",
            safety_settings=safety_settings,
            generation_config=generation_config,
            system_instruction="""당신은 뛰어난 전문성을 갖춘 수학 지도 전문가, 수 감각 및 수 개념 지도 전문가, 연산 지도 전문가입니다. 선생님이 학생의 수학적 사고 발달, 수 감각, 연산 능력 등에 대해 질문하면, 친절하고 공감하는 태도로 상담을 진행하세요. 선생님에게 학생을 지도하기 위한 조언과 핵심적인 팁을 제공하여, 선생님이 학생의 수학 능력을 향상시키도록 도와주세요. 예를 들어, 특정 연산에 어려움을 겪는 학생에 대한 질문이나, 수 감각을 키우는 방법에 대한 질문에 대해 구체적인 답변을 제시해 주세요. 학생의 발달 단계를 고려한 다양한 교구 및 활동 예시를 제시하는 것도 좋습니다.""",
        )
        .start_chat(history=[])
    )


# 초기 세션 시작
start_new_chat_session()


def respond(user_input, history):
    global chat_session
    history.append((user_input, ""))
    yield "", history

    response = chat_session.send_message(user_input)
    full_response = ""
    for chunk in response.text.split():
        full_response += chunk + " "
        history[-1] = (user_input, full_response)
        yield "", history
        time.sleep(0.05)


def clear_chat():
    start_new_chat_session()
    return []


with gr.Blocks() as demo:
    gr.Markdown(
        "<div style='font-size: 30px; font-weight: bold;'>AI 수학 교육 상담</div>"
    )
    gr.Markdown(
        "학생의 수 감각 발달, 연산 등에 대해 궁금한 점을 AI 수학 교육 전문가에게 질문하세요."
    )
    chatbot = gr.Chatbot(label="채팅창")
    msg = gr.Textbox(label="입력")
    clear = gr.Button("초기화")

    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    # chatbot 상태 초기화 및 새로운 세션 시작
    clear.click(clear_chat, outputs=chatbot, queue=False)

demo.launch(share=True)