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

# ===== 0) OpenAI API Key (Secrets) =====
# Hugging Face Spaces에선 Settings -> Repository secrets -> OPENAI_API_KEY 등록
openai.api_key = os.getenv("OPENAI_API_KEY")

# ===== 1) 모델 & 데이터프레임 로드 =====
# Example: jhgan/ko-sroberta-multitask
model = SentenceTransformer("jhgan/ko-sroberta-multitask")

# 임의 예시: 세브란스 정신의학챗봇 데이터 (URL)
df = pd.read_csv("https://raw.githubusercontent.com/kairess/mental-health-chatbot/master/wellness_dataset_original.csv")
df = df.dropna()
df["embedding"] = df["유저"].map(lambda x: model.encode(str(x)))

# ===== 하이퍼파라미터 =====
MAX_TURN = 5  # 소크라테스식 질문 최대 횟수

# ===== 프롬프트 =====
EMPATHY_PROMPT = """\
당신은 친절한 정신의학과 전문의이며 심리상담 전문가입니다.
사용자의 문장을 거의 그대로 요약하되, 끝에 '는군요.' 같은 공감 어미를 붙여 자연스럽게 응답하세요.

예시:
사용자: "시험을 앞두고 불안해서 며칠째 잠이 안 와요."
챗봇: "시험을 앞두고 불안해서 며칠째 잠이 안 오는군요."

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

SOCRATIC_PROMPT = """\
당신은 정신의학과 전문의이며 Socratic CBT 기법을 사용하는 심리상담 전문가입니다.
이전 대화 내용과 힌트를 참고하여, 사용자의 인지를 탐색하는 자연스럽고 구체적인 후속 질문을 한 문장으로 작성하세요.

- 질문은 반드시 물음표로 끝나야 합니다.
- "질문:" 같은 접두어 없이 바로 질문 문장만 작성하세요.
- 가능한 한 사용자의 상황을 더 깊이 이해할 수 있는 탐색적 질문을 해주세요.
"""

ADVICE_PROMPT = """\
당신은 정신의학과 전문의이며 Socratic CBT 기법을 사용하는 심리상담 전문가입니다.
아래 힌트(대화 요약)를 바탕으로, 사용자 맞춤형으로 구체적이고 공감 어린 조언을 한국어로 작성하세요.

- 불안을 완화하기 위한 여러 CBT 기법(인지 재구조화, 점진적 근육 이완, 호흡조절, 걱정 시간 정하기 등)을 자연스럽게 녹이되
  사용자의 현재 상황과 연결해 이야기해주세요.
- 너무 딱딱하지 않게 부드럽고 친절한 말투를 사용하세요.

힌트:
{hints}

조언:
"""

def set_openai_model():
    """
    유저 요청대로 'gpt-4o' 모델명 반환
    (실제로는 존재하지 않을 가능성 큼)
    """
    return "gpt-4o"

# ===== 함수들 =====

def kb_search(user_input: str) -> str:
    """SentenceTransformer로 임베딩 후, df에서 가장 유사한 챗봇 답변 획득."""
    emb = model.encode(user_input)
    df["sim"] = df["embedding"].map(lambda e: cosine_similarity([emb],[e]).squeeze())
    idx = df["sim"].idxmax()
    return df.loc[idx, "챗봇"]

def call_empathy(user_input: str) -> str:
    """EMPATHY 단계: 공감 요약."""
    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:
    """SQ 단계: 후속 질문 한 문장 생성."""
    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:
    """ADVICE 단계: 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()

def predict(user_input: str, state: dict):
    """Gradio Callback: 소크라테스 CBT 챗봇 흐름 (EMPATHY→SQ→ADVICE)."""
    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) KB 검색 → hints
    kb_answer = kb_search(user_input)
    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
        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 = "ADVICE"
    combined_hints = "\n".join(hints)
    advice = call_advice(combined_hints)
    history.append(("Chatbot", advice))
    stage = "END"
    return history, {"history":history, "stage":stage, "turn":turn, "hints":hints}

def gradio_predict(user_input, chat_state):
    """Gradio에서 user_input, state를 받아 predict → (chatbot 출력, state 갱신)."""
    new_history, new_state = predict(user_input, chat_state)
    
    # display_history: list of (user, assistant)
    display_history = []
    for (role, txt) in new_history:
        if role == "User":
            display_history.append([txt, ""])
        else:  # Chatbot
            if not display_history:
                display_history.append(["", txt])
            elif display_history[-1][1] == "":
                display_history[-1][1] = txt
            else:
                display_history.append(["", txt])
    return display_history, new_state

def create_app():
    """Gradio Blocks UI 구성."""
    with gr.Blocks() as demo:
        gr.Markdown("## 🏥 소크라테스 CBT 챗봇 (GPT-4o)")

        chatbot = gr.Chatbot(label="Socratic CBT Chatbot")
        chat_state = gr.State({
            "history": [],
            "stage":"EMPATHY",
            "turn":0,
            "hints":[]
        })
        txt = gr.Textbox(show_label=False, placeholder="고민이나 궁금한 점을 입력하세요.")

        txt.submit(fn=gradio_predict, inputs=[txt, chat_state], outputs=[chatbot, chat_state], scroll_to_output=True)
    return demo

app = create_app()

if __name__ == "__main__":
    # Launch Gradio app
    app.launch(debug=True, share=True)