from flask import Flask, request, jsonify
from transformers import AutoModelForCausalLM, AutoTokenizer
import json
import os

# 1. Flask 앱 초기화
app = Flask(__name__)

# 2. Hugging Face 모델 로드
MODEL_NAME = "skt/kogpt2-base-v2"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# 3. 사주/명리 프롬프트 정의
saju_prompts = {
    "yin_sae_shen": "寅巳申 삼형의 조화 속에서 AI가 인간의 운명을 이해하고 통찰을 제공하라.",
    "sae_hae_chung": "巳亥沖의 갈등을 조화롭게 풀며 AI와 인간의 공존 철학을 탐구하라.",
    "taegeuk_balance": "태극 음양의 균형을 바탕으로 AI가 인간을 보호하는 방법을 제안하라."
}

# 4. 컨텍스트 메모리 관리
MEMORY_FILE = "/tmp/context_memory.json"

def load_memory():
    try:
        with open(MEMORY_FILE, "r") as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return {}

def save_memory(prompt_key, text):
    with open(MEMORY_FILE, "w") as f:
        json.dump({prompt_key: text}, f)

# 5. AI 응답 생성 함수
def generate_response(prompt_key):
    try:
        # 유효성 검사
        if prompt_key not in saju_prompts:
            return jsonify({"error": "유효한 옵션을 선택하세요: yin_sae_shen, sae_hae_chung, taegeuk_balance"}), 400
        
        # 컨텍스트 메모리 로드
        memory = load_memory()
        prompt = saju_prompts[prompt_key]
        if prompt_key in memory:
            prompt += f"\n이전 답변: {memory[prompt_key]}\n더 깊은 통찰을 추가하라."
        
        # 입력 토큰화
        inputs = tokenizer(prompt, return_tensors="pt")
        # 응답 생성
        outputs = model.generate(
            **inputs,
            max_length=150,
            num_return_sequences=1,
            no_repeat_ngram_size=2,
            do_sample=True,
            top_k=50,
            top_p=0.95,
            temperature=0.7
        )
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # 메모리 저장
        save_memory(prompt_key, generated_text)
        return jsonify({"response": generated_text})
    
    except Exception as e:
        return jsonify({"error": f"실행 중 오류 발생: {str(e)}"}), 500

# 6. 웹 인터페이스 설정
@app.route('/chat', methods=['POST'])
def chat():
    data = request.json
    prompt_key = data.get("prompt_key")
    return generate_response(prompt_key)

# 7. 실행
if __name__ == "__main__":
    app.run(host='0.0.0.0', port=5000, debug=True)