#!/usr/bin/env python
# coding: utf-8

# ## 1. tsv full data load
import pandas as pd

df = pd.read_csv("sl_webtoon_full_data_sequential.tsv", sep="\t")

print(df.head())
print("전체 문장 수:", len(df))
print("컬럼 목록:", df.columns.tolist())

df['row_id'] = df.index  # 인덱스 컬럼 추가
df['text'] = df.apply(
    lambda x: f"[{x['에피소드']}] #{x['row_id']} {x['type']} {x['scene_text']}",
    axis=1
)
texts = df['text'].tolist()
print("최종 문장 수:", len(texts))

# ## 2. RAG 문장 생성
print("예시 5개:")
for t in df['text'].head(5).tolist():
    print("-", t)

# ## 3. 한국어 임베딩 모델 로드, 벡터 db
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings

embedding_model = HuggingFaceEmbeddings(model_name='jhgan/ko-sroberta-multitask')

db = FAISS.from_texts(texts, embedding_model)
print(" 벡터DB 생성 완료. 총 문장 수:", len(texts))
db.save_local("solo_leveling_faiss_ko")

db = FAISS.load_local("solo_leveling_faiss_ko", embedding_model, allow_dangerous_deserialization=True)

# 검색 테스트
query = "마나석이 뭐지?"
docs = db.similarity_search(query, k=5)
for i, doc in enumerate(docs, 1):
    print(f"[{i}] {doc.page_content}")

# ## 4. LLM 로드 (CPU 전용)
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain_community.llms import HuggingFacePipeline
import torch

# CPU로 강제
generator = pipeline(
    "text-generation",
    model="kakaocorp/kanana-nano-2.1b-instruct",
    device=-1   # ✅ CPU 사용
)

embedding_model = HuggingFaceEmbeddings(model_name='jhgan/ko-sroberta-multitask')
vectorstore = FAISS.load_local("solo_leveling_faiss_ko", embedding_model, allow_dangerous_deserialization=True)

model_name = "kakaocorp/kanana-nano-2.1b-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float32   # ✅ CPU에서는 float32
).to("cpu")                      # ✅ CPU 사용

llm_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=128)
llm = HuggingFacePipeline(pipeline=llm_pipeline)

custom_prompt = PromptTemplate(
    input_variables=["context", "question"],
    template="다음 문맥을 참고하여 질문에 답하세요.\n\n문맥:\n{context}\n\n질문:\n{question}\n\n답변:"
)

qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=vectorstore.as_retriever(search_kwargs={"k": 5}),
    chain_type="stuff",
    return_source_documents=True,
    chain_type_kwargs={"prompt": custom_prompt}
)

# 질문 테스트
query = "성진우는 몇 급 헌터지?"
result = qa_chain({"query": query})
print("답변:", result["result"])
print("\n참조 문서:")
for doc in result["source_documents"]:
    print(doc.page_content)

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

print("\n[선택지]")
for idx, choice in enumerate(choices, start=1):
    print(f"{idx}. {choice}")

user_idx = int(input("\n선택 번호 입력: ")) - 1
user_choice = choices[user_idx]
print(f"\n[사용자 선택]: {user_choice}")

result = qa_chain({"query": user_choice})
retrieved_context = "\n".join([doc.page_content for doc in result["source_documents"]])

print("\n[검색된 근거 문서 예시]")
print(retrieved_context[:600], "...") 

prompt = f"""
당신은 웹툰 '나 혼자만 레벨업'의 성진우입니다.
현재 상황:
{retrieved_context}
사용자 선택: {user_choice}
성진우의 말투로 간결하고 자연스러운 대사를 1~2문장 생성하세요.
중복된 내용이나 비슷한 문장은 만들지 마세요.
"""

response = generator(
    prompt, 
    max_new_tokens=200, 
    do_sample=True, 
    temperature=0.6,
    top_p=0.9,
    return_full_text=False
)[0]["generated_text"]

print("\n[성진우 응답]")
print(response)