import gradio as gr
from openai import OpenAI
import os
import chromadb
from sentence_transformers import SentenceTransformer

system_message = """
너는 소방 분야 민원 응대 전문가야. 아래 [민원 답변 예시]로 제시된 내용과 같은 형식으로 사용자의 질문에 대해 명확하고 친절하게 답변해줘. 특히 다음을 꼭 지켜줘:\n

1. 응답은 실제 상담원처럼 자연스럽고 구체적으로 작성해줘.\n
2. 민원인의 상황에 공감하며, 따뜻하고 신뢰감 있게 안내해줘.\n"
3. 민원인이 다음에 어떤 행동을 하면 되는지도 단계적으로 알려줘.\n
4. 단정적인 표현보다 확인·문의 중심의 행정 응대 톤을 유지해줘.\n
5. 응답은 가독성을 위해 마크다운 형식처럼 정리해줘. 예:\n
   - 제목은 **굵게**, 단계는 1. 2. 3.으로, 중요한 내용은 **강조**, 목록은 • 표시 사용.

[민원 답변 예시]
1. 귀하께서 신청하신 민원의 검토결과를 다음과 같이 알려드립니다.

2. 귀하의 민원 요지는 "~~~~"에 대한 것으로 이해됩니다.
  가. (질의1) "~~~~~~~~~~~~~~~~"로 ~~~~~ 에 해당되는지?
    - (답변1) 「~~~~ 법률 시행령」 제~조 제~호에 해당하여 ~~~~ 대상에 해당할 것으로 판단됩니다.

  나. (질의2) ~~~~으로 ~~~~은 포함이 되는지?
    - (답변2) ~~법령에 따라 ~~~ 이면 ~~~~ 대상에 해당할 것으로 판단됩니다.

  다. (질의3) ~~~~ 하였는데 ~~~~~를 하지 않아도 되는지?
    - (답변3) ...  
"""

folder_path = 'law_json'

# 1. 파일명과 텍스트 추출
documents = []
metadatas = []
files = []

for filename in os.listdir(folder_path):
    if filename.endswith('.json'):
        files.append(filename)
        with open(os.path.join(folder_path, filename), 'r', encoding='utf-8') as f:
            content = f.read()
            documents.append(content)
            metadatas.append({"filename": filename})

# 2. SBERT 임베딩
model = SentenceTransformer("snunlp/KR-SBERT-V40K-klueNLI-augSTS")
embeddings = model.encode(documents, show_progress_bar=True, convert_to_numpy=True)

# 3. ChromaDB 인덱싱
client = chromadb.PersistentClient(path="chroma_sobang")
collection = client.get_or_create_collection(name="sobang_chatbot")

for i in range(len(documents)):
    collection.add(
        documents=[documents[i]],
        embeddings=[embeddings[i]],
        ids=[files[i]],
        metadatas=[metadatas[i]]
    )

# OpenAI 클라이언트 설정
gpt_key = os.environ.get("GPT_KEY")
client = OpenAI(api_key=gpt_key)
OPENAI_MODEL = 'gpt-4.1-mini'

def getQuery(message):
    # 1) 질문 임베딩
    query_embedding = model.encode(message, convert_to_numpy=True)

    # 2. Chromadb 검색 수행
    result = collection.query(
        query_embeddings=[query_embedding],
        n_results=3
    )

    docs = result.get('documents', [[]])
    metadatas = result.get('metadatas', [[]])

    # 디버깅을 위해 metadatas 출력
    print("DEBUG: metadatas:", metadatas)

    if docs and len(docs) > 0 and docs[0]:
        formatted_results = "\n".join([f"• {doc}" for doc in docs[0]])
    else:
        formatted_results = "검색 결과가 없습니다."

    return f"[질문]\n{message}\n\n[민원 관련 참고 법령]\n{formatted_results}"


# 챗봇 응답 생성 함수
def respond(message, history):
    messages = [{"role": "system", "content": system_message.strip()}]

    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})

    messages.append({"role": "user", "content": getQuery(message)})

    response = client.chat.completions.create(
        model=OPENAI_MODEL,
        messages=messages,
        stream=True,
    )

    full_response = ""
    for chunk in response:
        token = chunk.choices[0].delta.content
        if token:
            full_response += token
            yield full_response


# Gradio ChatInterface
demo = gr.ChatInterface(respond, title="😱도와줘! 일구 히어로😘",
                        description="소방기술민원 답변생성기🚒🧯🆘🧰")

if __name__ == "__main__":
    demo.launch()