import os
import openai
import gradio as gr
import pandas as pd
import numpy as np
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

# OpenAI API Key (Hugging Face Secrets)
openai.api_key = os.getenv("OPENAI_API_KEY", "")

# =============== 0) 모델 / df 준비 ===============
# SentenceTransformer
model = SentenceTransformer('jhgan/ko-sroberta-multitask')

# 정신의학챗봇 CSV 로드
df = pd.read_csv('https://raw.githubusercontent.com/kairess/mental-health-chatbot/master/wellness_dataset_original.csv')
df = df.dropna()
# Unnamed 컬럼 제거
if 'Unnamed: 3' in df.columns:
    df = df.drop(columns=['Unnamed: 3'])

# 임베딩 필드
df['embedding'] = df['유저'].map(lambda x: model.encode(str(x)))

# ============== 1) 파라미터/프롬프트 ==============
MAX_TURN = 5  # 최대 소크라테스 질문 회수

def set_openai_model():
    """
    GPT-4 대신 'gpt-4o' (실제론 비존재 모델) 
    => 실제로는 'gpt-3.5-turbo' 등으로 교체 권장
    """
    return "gpt-4o"

EMPATHY_PROMPT = """\
당신은 친절한 정신의학과 전문의이며 심리상담 전문가입니다.

사용자의 문장을 거의 그대로 요약하되, 끝에 '는군요.' 같은 공감 어미로 자연스럽게 응답하고,
그 다음 줄에 한 문장으로 질문을 작성하세요.

유의사항:
1) 첫 문장은 공감형 요약 (예: "시험을 앞두고 불안해서 며칠째 잠을 못 자고 계시는군요.")
2) 두 번째 문장은 탐색/유도 질문 
   - 예: "어떤 고민들이 밤에 가장 많이 떠오르시나요?"

(예시)
사용자: "시험을 앞두고 불안해서 며칠째 잠이 안 와요."
챗봇:
"시험을 앞두고 불안해서 며칠째 잠을 못 자고 계시는군요.
시험 기간이 다가올 때 가장 힘드신 부분은 무엇인가요?"

이제 사용자 발화를 아래에 주겠습니다:
사용자 발화: "{sentence}"
챗봇:
"""

SOCRATIC_PROMPT = """\
당신은 정신의학과 전문의이며 Socratic CBT 기법을 사용하는 심리상담 전문가입니다.

아래 '대화 힌트'에는 사용자가 직전까지 이야기한 상황이나 고민이 요약되어 있다고 가정합니다.  
이 내용에 **공감**을 표시한 뒤, 그 흐름을 이어받아 **자연스럽고 구체적인 후속 질문**을 한 문장으로 작성하세요.

**세부 지침**:
1) 첫 문장은 사용자의 상황을 간단히 공감해 주되, 끝에 '는군요.' 등의 어미로 자연스럽게 마무리하세요.  
   - 예: "시험 기간 동안 정말 많은 부담을 느끼고 계시는군요."
2) 두 번째 문장은 탐색/유도 질문을 딱 한 문장으로 작성하세요.  
   - '질문:' 같은 접두어는 쓰지 말고, 바로 문장으로 시작합니다.
   - 반드시 물음표로 끝나야 합니다 (예: "...어떤 것들이 가장 힘드셨나요?").
3) 질문은 사용자의 현재 고민과 직접적으로 연결되어, 심층적인 자기 탐색을 유도해야 합니다.
   - 예: "밤에 들려오는 어떤 생각들이 잠을 더 설치게 만드는지 혹시 떠오르시나요?"
4) Bullet Point나 목록 대신, 간단히 두 줄(공감 + 질문) 구조로 작성하되, 너무 길게 쓰지 말고 부드러운 톤을 유지하세요.

(예시)
사용자 발화: "남편이 비트코인 투자로 속을 썩이네"
챗봇:
"남편분의 투자 문제로 속이 많이 상하시는군요.
혹시 그로 인해 가장 힘들다고 느끼는 부분은 무엇인가요?"

이제 아래 '대화 힌트'를 참조하여, 1줄 공감 + 1줄 질문 두 줄로 답변해 주세요.

대화 힌트:
{context}
"""

ADVICE_PROMPT = """\
당신은 정신의학과 전문의이며 Socratic CBT 기법을 사용하는 심리상담 전문가입니다.

아래 힌트(대화 요약)와 함께, 다음에 제시된 5가지 CBT 기법을 적절히 참고하여,
사용자 맞춤형으로 구체적이고 공감 어린 조언을 한국어로 작성하세요:

(1) 수면 제한 요법 (Sleep Restriction):
   "수면 제한 요법은 침대에 머무는 시간을 의도적으로 줄여, 침대와 수면 사이의 올바른 연결고리를 재구축하는 방법입니다.
    예를 들어, 침대에 10시간 머물지만 실제 수면 시간이 5시간인 경우, 처음에는 5시간만 침대에서 자고 점차 시간을 늘려가며
    '침대는 수면을 위한 장소'로 인식하도록 돕습니다."

(2) 자극 조절 요법 (Stimulus Control):
   "자극 조절 요법은 침대와 수면의 환경을 재정립하며, 침대를 오직 수면만을 위한 장소로 인식하게 만드는 치료법입니다.
    예를 들어, 침대에 누워 있을 때는 즉시 잠들지 못하더라도, 침대에서는 오직 수면을 취하는 습관을 기르는 것이 핵심입니다."

(3) 수면 위생 교육 (Sleep Hygiene):
   "수면 위생 교육은 건강한 수면을 위해 생활 습관을 개선하는 방법입니다.
    카페인 섭취를 줄이거나 늦은 시간의 전자기기 사용·밝은 조명 등을 피하고, 낮에는 가벼운 운동을 해두는 등의 습관을 포함합니다."

(4) 이완 기법 (Relaxation Techniques):
   "이완 기법은 심호흡, 점진적 근육이완, 명상 같은 방법을 통해 자연스러운 수면을 유도하는 방법입니다.
    몸을 스캔하고, 거북한 스트레칭을 풀고, 근육 이완을 연습하며 긴장을 낮추는 것이 주된 목표입니다."

(5) 인지 재구성 (Cognitive Restructuring):
   "인지 재구성은 ‘우리가 상황을 어떻게 바라보느냐에 따라 몸의 반응도 달라질 수 있다’는 긍정적 관점으로 전환시키며,
    걱정이나 불안, 부정적인 사고 패턴을 점검·조절하는 기법입니다. 이를 통해 사용자의 걱정을 완화하고
    자기효능감을 높이도록 돕습니다."

아래 사항을 꼭 반영해 주세요:
- 불안을 완화하기 위한 위 기법들을 자연스럽게 녹이되, 사용자의 현재 상황(힌트에 담긴 고민)과 연결해 이야기하세요.
- 너무 딱딱하지 않게, 부드럽고 친절한 말투를 사용하세요.

힌트:
{hints}

조언:
"""

# ============== 2) OpenAI 호출 함수들 ==============

def call_empathy(user_input: str) -> str:
    """ 공감 요약 생성 """
    prompt = EMPATHY_PROMPT.format(sentence=user_input)
    resp = openai.ChatCompletion.create(
        model=set_openai_model(),
        messages=[
            {"role":"system","content":"당신은 친절한 심리상담 전문가입니다."},
            {"role":"user","content":prompt}
        ],
        max_tokens=150,
        temperature=0.7
    )
    return resp.choices[0].message.content.strip()

def call_socratic_question(context: str) -> str:
    """ 소크라테스 후속질문 1문장 생성 """
    prompt = f"{SOCRATIC_PROMPT}\n\n대화 힌트:\n{context}"
    resp = openai.ChatCompletion.create(
        model=set_openai_model(),
        messages=[
            {"role":"system","content":"당신은 Socratic CBT 전문가입니다."},
            {"role":"user","content":prompt}
        ],
        max_tokens=200,
        temperature=0.7
    )
    return resp.choices[0].message.content.strip()

def call_advice(hints: str) -> str:
    """ 최종 CBT 조언 """
    final_prompt = ADVICE_PROMPT.format(hints=hints)
    resp = openai.ChatCompletion.create(
        model=set_openai_model(),
        messages=[
            {"role":"system","content":"당신은 Socratic CBT 기법 전문가입니다."},
            {"role":"user","content":final_prompt}
        ],
        max_tokens=700,
        temperature=0.8
    )
    return resp.choices[0].message.content.strip()

# ============== 3) predict 함수: EMPATHY→SQ→ADVICE ==============
def predict(user_input: str, state: dict):
    history = state.get("history", [])
    stage = state.get("stage", "EMPATHY")
    turn = state.get("turn", 0)
    hints = state.get("hints", [])

    # 1) 사용자 발화 기록
    history.append(("User", user_input))

    # 2) 유사도 계산 → df['챗봇']
    query_emb = model.encode(user_input)
    df["sim"] = df["embedding"].map(lambda emb: cosine_similarity([query_emb],[emb]).squeeze())

    # idxmax() 에러 방지: df가 비었거나 sim이 NaN인 경우 처리
    if df["sim"].count() == 0:
        # fallback: 그냥 "지식베이스가 비어 있습니다" 등
        kb_answer = "적합한 지식베이스 응답을 찾지 못했어요."
    else:
        kb_answer = df.loc[df["sim"].idxmax(), "챗봇"]

    hints.append(f"[KB] {kb_answer}")

    # 3) 단계별 분기
    if stage == "EMPATHY":
        empathic = call_empathy(user_input)
        history.append(("Chatbot", empathic))
        hints.append(empathic)
        stage = "SQ"
        turn = 0
        return history, {"history": history, "stage": stage, "turn": turn, "hints": hints}

    if stage == "SQ" and turn < MAX_TURN:
        # 전체 대화 + hints → 소크라테스 질문
        context_text = "\n".join([f"{r}: {c}" for (r,c) in history]) + "\n" + "\n".join(hints)
        sq = call_socratic_question(context_text)
        history.append(("Chatbot", sq))
        hints.append(sq)
        turn += 1
        return history, {"history": history, "stage": stage, "turn": turn, "hints": hints}

    # ADVICE 단계
    stage = "END"
    combined_hints = "\n".join(hints)
    advice = call_advice(combined_hints)
    history.append(("Chatbot", advice))

    return history, {"history":history, "stage":stage, "turn":turn, "hints":hints}

# ============== 4) Gradio UI ==============
def gradio_predict(user_input, chat_state):
    new_history, new_state = predict(user_input, chat_state)

    # display_history: list of [사용자문자열, 챗봇문자열]
    display_history = []
    for (role, txt) in new_history:
        if role == "User":
            display_history.append([txt, ""])
        else:
            if len(display_history) == 0:
                display_history.append(["", txt])
            else:
                display_history[-1][1] = txt

    # 세 번째 값으로 "" 반환하면, textbox가 자동으로 비워짐
    return display_history, new_state, ""

def create_app():
    with gr.Blocks() as demo:
        chatbot = gr.Chatbot()
        chat_state = gr.State({"history": [], "stage": "EMPATHY", "turn": 0, "hints": []})
        txt = gr.Textbox(show_label=False, placeholder="혹시 잠을 이루지 못하고 계신가요? 당신의 이야기를 듣고 싶어요! 걱정이 있으시면 편하게 말씀해주세요 :D")

        # outputs에 txt를 추가 + 콜백에서 세 번째 값을 ""로 리턴
        txt.submit(
            fn=gradio_predict,
            inputs=[txt, chat_state],
            outputs=[chatbot, chat_state, txt]
        )

    return demo

app = create_app()

if __name__ == "__main__":
    # 실제 배포/실행
    app.launch(debug=True, share=True)