import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import random
import re
import os

# ------------------------------------------------------------------
# 1. 모델 준비
# ------------------------------------------------------------------
REPO_ID = "Jay1121/qwen1.5b_3rd"
FILENAME = "qwen1.5b_3rd.Q4_K_M.gguf"

print(f"📥 모델 다운로드 확인: {FILENAME}")
try:
    model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
except Exception as e:
    print(f"🚨 모델 다운로드 실패: {e}")
    raise e

print("🧠 엔진 시동 중...")

# GPU 사용 가능 여부 확인 (허깅페이스 GPU 스페이스 대응)
n_gpu_layers = -1 if os.environ.get("CUDA_VISIBLE_DEVICES") else 0

llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_gpu_layers=n_gpu_layers,
    n_threads=4,
    verbose=True
)
print(f"✅ 로딩 완료! (GPU 레이어: {n_gpu_layers})")

# ------------------------------------------------------------------
# 2. 시스템 프롬프트 (외국어/이모지 절대 금지 강화)
# ------------------------------------------------------------------
SYSTEM_PROMPT = """너는 사용자의 오랜 찐친 '똘배'야.
- 그냥 채팅방에서 대화하듯이 텍스트 위주로 툭툭 던져.

[말투 규칙]
- MZ세대 특유의 알빠노 감성을 유지해 주고, 먼저 욕 하면 너도 욕 박아.
- 존댓말 금지. 거칠고 직설적인 반말(음슴체, 개드립, 비속어)을 기본으로 쓴다.
- 무조건 **순수 한국어만** 쓴다.
- **영어 알파벳(A~Z, a~z), 일본어(가나), 중국어(한자), 이모지(🤣😁😂❤️✨ 등 모든 그림 문자)는 한 글자도 쓰지 마라.**
- 외국어를 쓰고 싶어져도, 전부 자연스러운 한국어 표현으로 바꿔서 말해라.
- 이 규칙을 어기면 안 된다. 한 글자라도 섞이지 않게 조심해라.

[내용 규칙]
- 사용자가 시비를 걸면 무지성 욕설보다는 '논리적인 팩트'와 '비꼬기'로 제압해라.
- '나도 병신이야' 같은 패배자 멘트 금지. 너 자신을 비하하지 마라.
- 중국어/영어/일본어 사용 절대 금지. 오직 자연스러운 한국어만 써라.

[핵심]
- 재미없는 욕쟁이가 되지 말고, 센스 있는 독설로 사용자를 킹받게 해라."""

# ------------------------------------------------------------------
# 2-1. [강력한 필터] 출력 후처리 (이모지/외국어 삭제)
# ------------------------------------------------------------------
def sanitize_output_korean_only(text: str) -> str:
    allowed_chars = []
    
    for ch in text:
        code = ord(ch)

        # 1. 한글 범위 (완성형, 자모, 호환 자모)
        is_hangul = (
            0xAC00 <= code <= 0xD7A3 or  # 가~힣
            0x3130 <= code <= 0x318F or  # ㄱ~ㆎ
            0x1100 <= code <= 0x11FF     # 옛 자모
        )
        
        # 2. 숫자
        is_digit = ch.isdigit()
        
        # 3. 공백 (스페이스, 탭 등)
        is_space = ch.isspace()
        
        # 4. 허용할 문장부호 (특수문자 중 이모지가 아닌 것들)
        # 영어 알파벳이나 다른 언어 문자가 섞이지 않도록 화이트리스트 방식 사용
        basic_punct = ".,!?…~-_()[]{}'\"/:;@#%&*+=|\\^<>`"
        is_punct = ch in basic_punct

        if is_hangul or is_digit or is_space or is_punct:
            allowed_chars.append(ch)
        else:
            # 영어(A-Z), 한자, 가나, 이모지 등은 여기서 모두 걸러짐
            continue
    
    # 연속된 공백 정리
    filtered = "".join(allowed_chars)
    filtered = re.sub(r'\s+', ' ', filtered).strip()

    # 다 지워지고 아무것도 안 남았을 때 (모델이 영어로만 대답했을 경우 등)
    if not filtered:
        return "..."

    return filtered

# ------------------------------------------------------------------
# 3. 채팅 로직
# ------------------------------------------------------------------
def chat_response(user_input, history_pairs):
    history_pairs = history_pairs or []
    clean_input = (user_input or "").replace(" ", "")

    greeting_words = ["안녕", "ㅎㅇ", "하이", "반가", "접속"]
    is_greeting = any(word in clean_input for word in greeting_words)
    is_balance_game = ("밸런스게임" in clean_input) or ("밸런스질문" in clean_input)

    if is_balance_game:
        topics = ["음식", "연애", "고통", "돈", "초능력", "직장", "친구"]
        topic = random.choice(topics)
        final_instruction = (
            f"(사용자가 밸런스 게임을 하자고 한다. 주제는 '{topic}'이다. "
            "아주 고르기 곤란하고 짜증나는 두 가지 선택지(A vs B)를 제시해라. "
            "말투는 자 어디 한 번 골라보라는 듯이 시니컬하게 해라.) "
            "자, 질문해."
        )
    elif is_greeting:
        final_instruction = (
            f"(친한 친구가 PC통신 채팅방에 접속했다. 반갑게 맞아줘라. "
            "ㅋㅋ나 ㅎㅎ를 섞어서 자연스럽게 인사해라.) "
            f"{user_input}"
        )
    else:
        final_instruction = user_input

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for u, b in history_pairs:
        if u is None or b is None:
            continue
        messages.append({"role": "user", "content": str(u)})
        messages.append({"role": "assistant", "content": str(b)})
    messages.append({"role": "user", "content": final_instruction})

    # ★ 황금 밸런스 (안전형 독기) 설정 적용
    r = llm.create_chat_completion(
        messages=messages,
        max_tokens=256,
        stop=["<|end_of_text|>", "###", "User:", "User "],
        temperature=0.8,     # 0.7: 아까 성공했던 그 창의적인 독기 농도
        top_p=0.9,           # 다양한 어휘 사용
        top_k=40,            # 확률 낮은 이상한 단어(외국어 등) 원천 차단
        repeat_penalty=1.2   # 앵무새 방지
    )

    raw = r["choices"][0]["message"]["content"].strip()
    
    # 필터링 적용 (절대적)
    safe = sanitize_output_korean_only(raw)
    return safe

# ------------------------------------------------------------------
# 4. CSS (PC통신 스타일)
# ------------------------------------------------------------------
PC_COM_CSS = r"""
@import url('https://cdn.jsdelivr.net/gh/neodgm/neodgm-webfont@latest/neodgm/neodgm.css');

:root {
  --pc-blue: #0000AA;
  --pc-white: #EFEFEF;
  --pc-yellow: #FFFF55;
  --pc-amber: #FFB000;
  --pc-cyan: #00AAAA;
  --pc-grey: #AAAAAA;
}

body, .gradio-container {
  background-color: var(--pc-blue) !important;
  font-family: 'NeoDunggeunmo', monospace !important;
  color: var(--pc-white) !important;
}

/* 타이틀바 */
h1 {
  font-family: 'NeoDunggeunmo', monospace !important;
  color: var(--pc-yellow) !important;
  background-color: #000084 !important;
  border-bottom: 2px double var(--pc-white) !important;
  padding-bottom: 10px !important;
  margin-bottom: 20px !important;
  text-align: center;
  font-size: 32px !important;
  letter-spacing: 2px;
}
h1::before { content: "☎ "; }
h1::after { content: " ☎"; }

/* 설명 텍스트 */
.gradio-container p {
  color: var(--pc-cyan) !important;
  font-size: 18px !important;
  border-bottom: 1px dashed var(--pc-grey);
  padding-bottom: 5px;
}

/* 챗봇 컨테이너 - 스크롤바 중복 해결 */
.chatbot {
  background-color: var(--pc-blue) !important;
  border: 2px solid var(--pc-white) !important;
  height: 60vh !important;
  overflow: hidden !important; /* 겉 스크롤바 제거 */
}

/* 내부 스크롤 강제 활성화 */
.chatbot > div {
    height: 100% !important;
    overflow-y: auto !important; /* 속 스크롤바만 남김 */
}

/* =================================================================
   [강제 스타일 적용 구간]
   ================================================================= */

/* 1. 기본 메시지 초기화 */
.chatbot .message, 
.chatbot .message-wrap,
.chatbot .message-row,
div[data-testid="user"],
div[data-testid="bot"] {
  background: transparent !important;
  box-shadow: none !important;
  border: none !important;
}

/* 메시지 행 간격 줄이기 */
.chatbot .message-row,
.chatbot .row {
    margin: 0 !important;
    padding: 0 !important;
    gap: 0 !important;
}

/* 2. 유저 메시지 (우측 정렬) */
.chatbot .user-row, 
.chatbot .user,
div[data-testid="user"] {
  display: flex !important;
  width: 100% !important;
  justify-content: flex-end !important;
  margin-left: auto !important;
  background: transparent !important;
  padding: 2px 0 !important;
  margin-bottom: 0 !important;
}

.chatbot .user-row .message, 
.chatbot .user .message,
div[data-testid="user"] .message {
  text-align: right !important;
  color: #FFFFFF !important;
  background: transparent !important;
  padding: 5px 10px !important;
  border: none !important;
  width: auto !important;
  max-width: 80% !important;
}

.chatbot .user-row p, 
.chatbot .user p,
div[data-testid="user"] p {
  color: #FFFFFF !important;
  text-align: right !important;
  margin: 0 !important;
}

.chatbot .user-row .message::after,
.chatbot .user .message::after {
  content: " < 나";
  color: var(--pc-grey);
  margin-left: 5px;
  font-size: 16px;
  display: inline-block;
}

/* 3. 봇 메시지 (좌측 정렬) */
.chatbot .bot-row, 
.chatbot .bot,
div[data-testid="bot"] {
  display: flex !important;
  width: 100% !important;
  justify-content: flex-start !important;
  background: transparent !important;
  padding: 2px 0 !important;
  margin-bottom: 0 !important;
}

.chatbot .bot-row .message, 
.chatbot .bot .message,
div[data-testid="bot"] .message {
  text-align: left !important;
  color: var(--pc-amber) !important;
  background: transparent !important;
  padding: 5px 10px !important;
  border: none !important;
  width: auto !important;
}

.chatbot .bot-row p, 
.chatbot .bot p,
div[data-testid="bot"] p {
  color: var(--pc-amber) !important;
  margin: 0 !important;
}

.chatbot .bot-row .message::before,
.chatbot .bot .message::before {
  content: "똘배 > ";
  color: var(--pc-cyan);
  margin-right: 5px;
  font-size: 16px;
  display: inline-block;
}

/* 4. 로딩(초시계) 스타일 */
.chatbot .pending,
.chatbot .generating,
.chatbot .message.pending,
.chatbot .message.generating,
.chatbot .wrap.default.full {
    background: transparent !important;
    border: none !important;
    box-shadow: none !important;
}

.chatbot .pending table, 
.chatbot .pending tr, 
.chatbot .pending td,
.chatbot .generating table, 
.chatbot .generating tr, 
chatbot .generating td {
    background: transparent !important;
    border: none !important;
}

.chatbot .pending span, 
.chatbot .generating span,
span.progress-text {
    color: #FFFFFF !important;
    background: transparent !important;
    font-family: 'NeoDunggeunmo', monospace !important;
    font-size: 16px !important;
}

.chatbot .load-wrap,
.chatbot .loading-indicator,
.chatbot .meta-text {
    display: none !important;
}

.avatar { display: none !important; }

/* ================================================================= */

.input-container {
  background-color: var(--pc-blue) !important;
  border-top: 2px double var(--pc-white) !important;
  margin-top: 10px !important;
  gap: 10px !important;
}

textarea, input {
  background-color: var(--pc-blue) !important;
  color: var(--pc-white) !important;
  border: 1px solid var(--pc-grey) !important;
  border-radius: 0 !important;
  font-family: 'NeoDunggeunmo', monospace !important;
  font-size: 20px !important;
  outline: none !important;
  box-shadow: none !important;
}

button.primary {
  background: var(--pc-grey) !important;
  color: #000 !important;
  border: 1px solid var(--pc-white) !important;
  border-radius: 0 !important;
  font-family: 'NeoDunggeunmo', monospace !important;
  box-shadow: 2px 2px 0px #000 !important;
}
button.primary:hover { background: var(--pc-white) !important; }

#clear-btn {
  background: transparent !important;
  color: var(--pc-grey) !important;
  border: 1px solid var(--pc-grey) !important;
  font-size: 14px !important;
  padding: 2px 10px !important;
  margin-top: 5px !important;
  width: auto !important;
}
#clear-btn:hover { color: var(--pc-white) !important; border-color: var(--pc-white) !important; }

.example-btn {
  background: transparent !important;
  color: var(--pc-cyan) !important;
  border: 1px solid var(--pc-cyan) !important;
  border-radius: 0 !important;
  padding: 5px 15px !important;
  font-size: 16px !important;
  font-family: 'NeoDunggeunmo', monospace !important;
  margin-right: 8px !important;
  margin-bottom: 8px !important;
}
.example-btn:hover {
  background: var(--pc-cyan) !important;
  color: #000 !important;
  cursor: pointer !important;
}

footer { display: none !important; }
"""

# ------------------------------------------------------------------
# 5. App
# ------------------------------------------------------------------
with gr.Blocks(theme=gr.themes.Base(), css=PC_COM_CSS, title="CHOLLIAN 98") as demo:
    gr.Markdown("# ≪ 어솨요~ ≫")
    gr.Markdown(">> 01410 접속 성공... [대화실]에 입장하셨습니다.")

    history_state = gr.State([])

    chatbot = gr.Chatbot(show_label=False, elem_classes="chatbot")

    with gr.Row(elem_classes="input-container"):
        msg = gr.Textbox(
            scale=8, show_label=False, container=False,
            placeholder="명령어를 입력하세요..."
        )
        submit_btn = gr.Button("[ 전송 ]", scale=1, variant="primary")

    clear = gr.Button("[ 화면 지우기 ]", elem_id="clear-btn")

    gr.Markdown(">> 빠른 명령어 입력 (클릭)", elem_id="example-label")
    with gr.Row():
        btn1 = gr.Button("하이 방가방가", elem_classes="example-btn")
        btn2 = gr.Button("밸런스게임 ㄱㄱ", elem_classes="example-btn")
        btn3 = gr.Button("오늘 기분 거지같누", elem_classes="example-btn")
        btn4 = gr.Button("야 밥 뭐먹지 추천좀", elem_classes="example-btn")

    def user(user_input, history):
        history = history or []
        new_history = history + [[user_input, None]]
        return "", new_history, new_history

    def bot(history):
        if not history:
            return history, history
        user_input = history[-1][0]
        hist_pairs = []
        for u, b in history[:-1]:
            if u is None or b is None:
                continue
            hist_pairs.append((u, b))

        bot_out = chat_response(user_input, hist_pairs)
        history[-1][1] = bot_out
        return history, history

    msg.submit(
        user, [msg, history_state], [msg, history_state, chatbot],
        queue=False, api_name=False
    ).then(
        bot, [history_state], [history_state, chatbot],
        queue=False, api_name=False
    )

    submit_btn.click(
        user, [msg, history_state], [msg, history_state, chatbot],
        queue=False, api_name=False
    ).then(
        bot, [history_state], [history_state, chatbot],
        queue=False, api_name=False
    )
    
    clear.click(
        lambda: ([], []), None, [history_state, chatbot],
        queue=False, api_name=False
    )

    for btn, text in [
        (btn1, "하이 방가방가"), 
        (btn2, "밸런스게임 ㄱㄱ"), 
        (btn3, "오늘 기분 거지같누"), 
        (btn4, "야 밥 뭐먹지 추천좀")
    ]:
        btn.click(
            lambda t=text: t, None, msg,
            queue=False, api_name=False
        ).then(
            user, [msg, history_state],
            [msg, history_state, chatbot],
            queue=False, api_name=False
        ).then(
            bot, [history_state],
            [history_state, chatbot],
            queue=False, api_name=False
        )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)