import pandas as pd
import torch
from flask import Flask, request, Response, render_template_string
from transformers import AutoTokenizer, GPT2LMHeadModel
from dicttoxml import dicttoxml
import re
import traceback

app = Flask(__name__)

# --- hCaptcha 설정 관련 코드 전부 제거됨 ---

# 1. 모델 로드 (기존과 동일)
print("토크나이저 로딩 중...")
tokenizer = AutoTokenizer.from_pretrained("skt/kogpt2-base-v2", trust_remote_code=True)
print("모델 로딩 중...")
model = GPT2LMHeadModel.from_pretrained("skt/kogpt2-base-v2", trust_remote_code=True)
print("모델 로딩 완료!")

# 2. 데이터셋 로드 (기존과 동일)
try:
    df = pd.read_excel('dataset.xlsx')
    knowledge_list = df['데이터셋에 넣을 내용(*)'].tolist()
except Exception as e:
    print(f"데이터셋 로드 에러: {e}")
    knowledge_list = []

def find_relevant_context(query, top_n=2):
    """질문과 관련된 지식데이터 문장 최대 top_n개 찾아서 반환 (기존과 동일)"""
    query_words = query.replace(" ", "").lower()
    relevant_sentences = []
    for s in knowledge_list:
        s_text = str(s).replace(" ", "").replace("\n", "").lower()
        if any(word.replace(" ", "") in s_text for word in query.split()):
            relevant_sentences.append(s)
    if relevant_sentences:
        return " ".join(str(s) for s in relevant_sentences[:top_n])
    return ""

def ask_sayknow(query):
    try:
        context = find_relevant_context(query)
        persona_guide = (
            "너는 지식 기반 한국어 챗봇 Sayknow야. 자기소개(이름, 정체, 인사 등) 질문은 '저는 Sayknow입니다.'라고 답해. "
            "그 외엔 아래 참고해서 정확하고 자연스러운 한국어 문장으로 80자 이내로 답해.\n"
            "예시: Q: 분수의 덧셈이 뭐야?\nA: 분모가 같을 때 분자끼리 더하면 됩니다.\n"
        )
        info = context if context else "정보 없음"
        prompt = f"{persona_guide}---\n[정보]\n{info}\n[질문]\n{query}\n[답변] "

        # 이전 답변 로직 개선 (attention_mask 추가) - 이 부분은 잘 작동하고 있을 거야!
        tokenizer.pad_token = tokenizer.eos_token
        encoded_input = tokenizer.encode_plus(
            prompt,
            return_tensors='pt',
            truncation=True,
            padding=True
        )
        input_ids = encoded_input['input_ids']
        attention_mask = encoded_input['attention_mask']
        
        model.eval()
        with torch.no_grad():
            gen_ids = model.generate(
                input_ids,
                attention_mask=attention_mask,
                max_new_tokens=512,  # 답변이 잘리는 문제 방지를 위해 조금 늘려봤어! (60 -> 80)
                min_length=5,
                repetition_penalty=1.3,
                do_sample=True,
                top_k=30,
                top_p=0.85,
                pad_token_id=tokenizer.pad_token_id,
                temperature=0.5,
                num_beams=1
            )
        raw_response = tokenizer.decode(gen_ids[0], skip_special_tokens=True) # 원본 응답 저장
        
        # --- 응답 처리 로직 개선 버전 (index out of range 에러 방지) ---
        # 1. 모델이 생성한 전체 텍스트에서 프롬프트 부분 자르기 (반복되는 문제 방지)
        #    prompt가 raw_response의 시작 부분에 있다면 그 부분을 잘라낼게.
        if raw_response.startswith(prompt):
            extracted_answer = raw_response[len(prompt):].strip()
        else:
            extracted_answer = raw_response.strip()

        # 2. '답변:' 키워드를 기준으로 진짜 답변 부분 추출
        if "답변:" in extracted_answer:
            answer = extracted_answer.split("답변:", 1)[1].strip() # 첫 번째 "답변:" 이후만
        else:
            # 만약 "답변:" 태그가 없으면, 프롬프트의 지시사항 중복 등을 제거 시도
            persona_end_marker = "답해.\n" # persona_guide의 특정 끝 부분을 표시
            if persona_end_marker in extracted_answer:
                try:
                    answer = extracted_answer[extracted_answer.rindex(persona_end_marker) + len(persona_end_marker):].strip()
                except ValueError:
                    answer = extracted_answer # 안되면 그냥 전체 사용
            else:
                answer = extracted_answer # 그것도 없으면 그냥 전체 사용

        # 그래도 답변이 비어있으면 오류 메시지를 대체
        if not answer:
            answer = "죄송합니다. 질문에 대한 답변을 찾을 수 없거나 내용이 명확하지 않습니다."


        # 1. 의미 없는 수식/영문/특수문자/반복문자 등 필터링 (기존과 동일)
        # 이 부분을 먼저 한번 적용해서 answer가 엉뚱한 문자열이 되는 걸 방지
        answer = re.sub(r"[^가-힣0-9 .,!?~\n]", "", answer)
        answer = re.sub(r"([.,!?~])\1{2,}", r"\1", answer)
        answer = re.sub(r"[a-zA-Z]+", "", answer)
        answer = re.sub(r"[=^*/\\]+", "", answer)
        answer = re.sub(r"\s+", " ", answer).strip()

        # 2. 80자 이내로 자르기 (한글 기준) (기존과 동일)
        def truncate_korean(text, max_len=80):
            count = 0
            result = ""
            for ch in text:
                result += ch
                count += 1
                if count >= max_len:
                    break
            return result
        answer = truncate_korean(answer, 80)

        # 3. 문장 끝이 자연스럽지 않으면 마침표 추가
        if answer and answer[-1] not in ".!?": 
            answer += "."
        elif not answer: # 빈 문자열인데 '.' 찍으면 에러나니 한번 더 체크
            answer = "알 수 없는 오류가 발생했습니다." # 최후의 보루

        return answer
    except Exception as e:
        print(f"ask_sayknow 에러: {e}")
        traceback.print_exc()
        return f"내부 오류: {str(e)}" # 외부 사용자에게 보이는 메시지!

# 3. API (XML 응답 유지) (기존과 동일)
@app.route('/chatapi.html', methods=['GET'])
@app.route('/index.html', methods=['GET'])
def chat_api():
    query = request.args.get('askdata', '')
    if not query:
        result = {"status": "error", "message": "No data"}
    else:
        try:
            answer = ask_sayknow(query)
            result = {
                "service": "Sayknow",
                "question": query,
                "answer": answer
            }
        except Exception as e:
            print(f"chat_api 에러: {e}")
            traceback.print_exc()
            result = {
                "service": "Sayknow",
                "question": query,
                "answer": f"에러 발생: {str(e)}",
                "error": str(e)
            }

    xml_output = dicttoxml(result, custom_root='SayknowAPI', attr_type=False)
    return Response(xml_output, mimetype='text/xml')

# 4. 웹 UI (간단한 질문 폼 + 답변) - hCaptcha 코드 전부 제거!
@app.route('/', methods=['GET', 'POST'])
def index():
    answer = ""
    question = ""
    # error_message 제거

    if request.method == "POST":
        question = request.form.get('question', '')
        # hcaptcha_response 관련 로직 제거
        
        # hCaptcha 검증 로직 제거
        if question: # 질문이 있으면 바로 답변 생성!
            answer = ask_sayknow(question)

    html = f"""
    <html>
    <head>
        <title>Sayknow 챗봇</title>
        <!-- hCaptcha 스크립트 제거 -->
    </head>
    <body>
        <h2>Sayknow 한국어 챗봇</h2>
        <form method="post" action="/">
            <input type="text" name="question" value="{question}" placeholder="질문을 입력하세요" style="width:300px;" autofocus />
            <br/><br/>
            <!-- hCaptcha 위젯 제거 -->
            <!-- 에러 메시지 보여주는 부분 제거 -->
            <br/>
            <input type="submit" value="질문하기" />
        </form>
        <hr>
        <h3>답변:</h3>
        <p style="white-space: pre-wrap;">{answer}</p>
    </body>
    </html>
    """
    return render_template_string(html)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)