from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
from huggingface_hub import InferenceClient
import os
from dotenv import load_dotenv

load_dotenv()

app = Flask(__name__)
CORS(app, resources={
    r"/*": {
        "origins": ["http://your-frontend-domain"],
        "methods": ["GET", "POST", "OPTIONS"],
        "allow_headers": ["Content-Type", "Accept"]
    }
})

# Hugging Face 클라이언트 초기화
client = InferenceClient(
    model=os.getenv("MODEL_ID"),
    token=os.getenv("HUGGINGFACE_API_KEY")
)

# 더 큰 한국어 모델로 변경
model_name = "nlpai-lab/kullm-polyglot-5.8b-v2"  # 더 나은 품질의 모델

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True
)

# device_map="auto"를 사용하므로 .to(device) 호출 제거
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# model = model.to(device)  # 이 줄 제거

# 루트 경로 핸들러 추가
@app.route('/')
def serve_frontend():
    return send_from_directory('../frontend', 'index.html')

# 정적 파일 서빙을 위한 라우트 추가
@app.route('/<path:path>')
def serve_static(path):
    return send_from_directory('../frontend', path)

# favicon.ico 핸들러 추가
@app.route('/favicon.ico')
def favicon():
    return '', 204  # No Content 응답 반환

@app.route('/api/generate-diary', methods=['POST'])
def generate_diary():
    try:
        data = request.json
        if not data or 'keywords' not in data:
            return jsonify({"error": "키워드가 필요합니다"}), 400
            
        keywords = data.get('keywords', '').strip()
        if not keywords:
            return jsonify({"error": "키워드가 비어있습니다"}), 400
        
        prompt = f"""다음은 오늘 있었던 일의 요약입니다. 이것을 바탕으로 생생하고 감동적인 일기를 작성해주세요.

        [상세 요구사항]
        1. 도입부: 
           - 그날의 날씨나 분위기로 시작
           - 상황과 등장인물 소개
        
        2. 전개:
           - 구체적인 대화와 행동 묘사
           - 오감을 사용한 장면 묘사
           - 등장인물들의 표정과 감정 변화
        
        3. 감정과 생각:
           - 내면의 감정을 섬세하게 표현
           - 사건에 대한 나의 생각과 깨달음
           - 다른 사람들의 감정에 대한 공감
        
        4. 문체:
           - 문어체와 구어체를 적절히 혼용
           - 비유와 은유를 활용한 표현
           - 반복을 피하고 다양한 어휘 사용
        
        5. 마무리:
           - 그날의 경험이 주는 의미
           - 앞으로의 기대나 다짐
        
        요약:
        {keywords}
        
        ===
        오늘의 일기:
        오늘은 """
        
        # Hugging Face API를 통한 텍스트 생성
        parameters = {
            "max_new_tokens": 768,
            "temperature": 0.88,
            "top_p": 0.95,
            "repetition_penalty": 1.35,
            "top_k": 50,
            "do_sample": True,
            "num_return_sequences": 1
        }
        
        response = client.text_generation(
            prompt,
            **parameters
        )
        
        if not response:
            return jsonify({"error": "일기 생성에 실패했습니다"}), 500
            
        # 프롬프트 제거하고 생성된 텍스트만 반환
        diary_content = response.split("오늘은 ")[-1].strip()
        diary_content = "오늘은 " + diary_content
            
        return jsonify({"diary": diary_content})
    
    except Exception as e:
        print(f"Error generating diary: {str(e)}")
        return jsonify({"error": f"일기 생성 중 오류가 발생했습니다: {str(e)}"}), 500

if __name__ == '__main__':
    app.run(debug=True)