import os, zipfile, shutil, glob
import gradio as gr
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch
import langchain

ZIP_NAME = "solo_leveling_faiss_ko.zip"
TARGET_DIR = "solo_leveling_faiss_ko"

def ensure_faiss_dir() -> str:
    if os.path.exists(os.path.join(TARGET_DIR, "index.faiss")) and \
       os.path.exists(os.path.join(TARGET_DIR, "index.pkl")):
        return TARGET_DIR
    if os.path.exists("index.faiss") and os.path.exists("index.pkl"):
        os.makedirs(TARGET_DIR, exist_ok=True)
        if not os.path.exists(os.path.join(TARGET_DIR, "index.faiss")):
            shutil.move("index.faiss", os.path.join(TARGET_DIR, "index.faiss"))
        if not os.path.exists(os.path.join(TARGET_DIR, "index.pkl")):
            shutil.move("index.pkl", os.path.join(TARGET_DIR, "index.pkl"))
        return TARGET_DIR
    if os.path.exists(ZIP_NAME):
        with zipfile.ZipFile(ZIP_NAME, 'r') as z:
            z.extractall(".")
        if os.path.exists(os.path.join(TARGET_DIR, "index.faiss")) and \
           os.path.exists(os.path.join(TARGET_DIR, "index.pkl")):
            return TARGET_DIR
        faiss_cand = glob.glob("**/index.faiss", recursive=True)
        pkl_cand   = glob.glob("**/index.pkl",   recursive=True)
        if faiss_cand and pkl_cand:
            os.makedirs(TARGET_DIR, exist_ok=True)
            shutil.copy2(faiss_cand[0], os.path.join(TARGET_DIR, "index.faiss"))
            shutil.copy2(pkl_cand[0],   os.path.join(TARGET_DIR, "index.pkl"))
            return TARGET_DIR
    raise FileNotFoundError("FAISS index files not found (index.faiss / index.pkl).")

# 0) FAISS 인덱스 위치 확보
base_dir = ensure_faiss_dir()

# 1) 벡터 DB
embeddings = HuggingFaceEmbeddings(model_name="jhgan/ko-sroberta-multitask")
vectorstore = FAISS.load_local(base_dir, embeddings, allow_dangerous_deserialization=True)

# 2) 모델 로딩 (CPU)
model_name = "kakaocorp/kanana-nano-2.1b-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32, device_map=None)

# 3) 텍스트 생성 파이프라인
lm = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=100,
    temperature=0.6,
    do_sample=True,
    top_p=0.9,
    return_full_text=False
)

# 선택지
choices = [
    "1: 황동석 무리를 모두 처치한다.",
    "2: 진호를 포함한 황동석 무리를 모두 처치한다.",
    "3: 전부 기절 시키고 살려둔다.",
    "4: 시스템을 거부하고 그냥 도망친다."
]

def rag_answer(message, history):
    try:
        user_idx = int(message.strip()) - 1
        user_choice = choices[user_idx]
    except:
        return "❗올바른 번호를 입력해주세요. (예: 1 ~ 4)"
    docs = vectorstore.similarity_search(user_choice, k=3)
    context = "\n".join([doc.page_content for doc in docs])
    prompt = f"""당신은 웹툰 '나 혼자만 레벨업'의 성진우입니다.
현재 상황:
{context}
사용자 선택: {user_choice}
성진우의 말투로 간결하고 자연스러운 대사를 1~2문장 생성하세요.
중복된 내용이나 비슷한 문장은 만들지 마세요.
"""
    response = lm(prompt)[0]["generated_text"]
    only_dialogue = response.strip().split("\n")[-1]
    if not only_dialogue.startswith("대사:"):
        only_dialogue = "대사: " + only_dialogue
    return only_dialogue

# ===== UI (변경 지점) =====
css_code = """
.quest-title {
  display:flex; align-items:center; gap:10px;
  font-weight:700; font-size:22px; margin-bottom:6px;
}
.quest-title img {
  width:72px; height:auto; opacity:.95;
}
.quest-desc { line-height:1.5; margin-bottom:14px; }
"""

header_html = """
<div class="quest-title">
  [긴급 퀘스트: 적을 처치하라!]
</div>
<div class="quest-desc">
  
<div style="display: flex; align-items: flex-start; gap: 20px;">
   
        <div class="quest-desc">
            '플레이어'에게 살의를 가진 이들이 주변에 있습니다...
            <!-- 나머지 텍스트 -->
        </div>
    </div>
    <div>
        <img src="https://huggingface.co/spaces/min24ss/r-story-selection/resolve/main/system.png"
             alt="quest"
             style="width: 250px; height: auto;">
    </div>
</div>
     

  '플레이어'에게 살의를 가진 이들이 주위에 있습니다. 이들을 모두 처치하여 안전을 확보하십시오.<br>
  지시에 따르지 않으면 당신의 심장은 정지(!)하게 됩니다.<br>
  처치해야 할 적의 숫자: 8명 / 처치한 적의 숫자: 0명<br><br>
  💬 선택지를 입력하세요:<br>
  1: 황동석 무리를 모두 처치한다.<br>
  2: 황동석 무리와 진호를 포함하여 모두 처치한다.<br>
  3: 전부 기절 시키고 살려둔다.<br>
  4: 시스템을 거부하고 그냥 도망친다.
</div>
"""

with gr.Blocks(css=css_code) as demo:
    gr.HTML(header_html)  # ← 여기서 HTML 그대로 렌더링 (이미지 보장)
    gr.ChatInterface(fn=rag_answer)  # title/description은 쓰지 않음

if __name__ == "__main__":
    print("Torch:", torch.__version__)
    import transformers as _t
    print("Transformers:", _t.__version__)
    print("LangChain:", langchain.__version__)
    demo.launch()