import os
import asyncio
from typing import List, Dict

# Protobuf C-extension 대신 pure-Python 구현 사용
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"

# .env 파일 및 Space Secrets 로드
from dotenv import load_dotenv
load_dotenv()

# Gradio client 버그 우회 (OpenAPI 파싱)
import gradio_client.utils as client_utils
orig = client_utils.json_schema_to_python_type
def safe_json_type(schema, defs=None):
    try:
        return orig(schema, defs)
    except Exception:
        return "Any"
client_utils.json_schema_to_python_type = safe_json_type

# Google API Key 검증
api_key = os.getenv("GOOGLE_API_KEY")
if not api_key:
    raise EnvironmentError("GOOGLE_API_KEY를 Settings→Secrets에 추가해주세요.")
os.environ["GOOGLE_API_KEY"] = api_key

# ChromaDB 경로 설정
db_dir = os.path.join(os.getcwd(), "chromadb_KH_media")
os.environ["CHROMA_DB_DIR"] = db_dir

# === 라이브러리 임포트 ===
import chromadb
import gradio as gr
from sentence_transformers import SentenceTransformer
from google.adk.agents import Agent
from google.adk.sessions import InMemorySessionService
from google.adk.runners import Runner
from google.genai import types

# === 현재 교수진 목록 === 
PROFESSORS = [
    "이인희", "김태용", "박종민", "홍지아", "이정교",
    "이기형", "이선영", "조수영", "이종혁", "이두황",
    "이상원", "이훈", "최수진", "최민아", "김관호"
]

# === Simple RAG 시스템 ===
class SimpleRAGSystem:
    def __init__(self, db_path: str = None, collection: str = "KH_media_docs"):
        db_path = db_path or os.getenv("CHROMA_DB_DIR")
        self.embedding_model = SentenceTransformer("snunlp/KR-SBERT-V40K-klueNLI-augSTS")
        self.client = chromadb.PersistentClient(path=db_path)
        self.collection = self.client.get_collection(name=collection)
        count = self.collection.count()
        if count == 0:
            raise RuntimeError("ChromaDB가 비어있습니다.")

    def search_similar_docs(self, query: str, top_k: int = 20) -> List[Dict]:
        emb = self.embedding_model.encode(query).tolist()
        res = self.collection.query(
            query_embeddings=[emb], n_results=top_k,
            include=["documents", "metadatas"]
        )
        docs = []
        for doc, meta in zip(res["documents"][0], res["metadatas"][0]):
            docs.append({"role": "system", "content": doc})
        return docs

rag_system = SimpleRAGSystem()

# === Google ADK 설정 ===
session_svc = InMemorySessionService()
agent = Agent(model="gemini-2.0-flash-lite", #"gemini-2.0-flash"
              name="khu_media_advisor",
              instruction="""당신은 경희대학교 미디어학과 전문 상담 AI입니다.
# 주요 역할:
- 제공된 문서 정보를 바탕으로 답변 제공
- 미디어학과 관련 질문에 친절하고 구체적으로 응답
- 문서에 없는 내용은 일반 지식으로 보완 (단, 명시)
# 답변 스타일:
- 자세하고 풍부한 설명을 포함하여 상세하고 길게 답변 제공
- 친근하고 도움이 되는 상담사 톤
- 핵심 정보를 명확하게 전달
- 추가 궁금한 점이 있으면 언제든 물어보라고 안내
# 참고 문서 활용:
- 문서 내용이 있으면 구체적으로 인용
- 여러 문서의 정보를 종합하여 답변 작성
- 정확하지 않은 정보는 추측하지 말고 솔직하게 모른다고 답변

# 현재 경희대학교 미디어학과 교수진은 다음과 같습니다:
이인희, 김태용, 박종민, 홍지아, 이정교, 이기형, 이선영, 조수영, 이종혁, 이두황, 이상원, 이훈, 최수진, 최민아, 김관호"""             
)

runner = Runner(agent=agent, app_name="khu_media_chatbot", session_service=session_svc)
session_id = None

async def get_response(prompt: str) -> str:
    global session_id
    if session_id is None:
        sess = await session_svc.create_session(app_name="khu_media_chatbot", user_id="user")
        session_id = sess.id
    content = types.Content(role="user", parts=[types.Part(text=prompt)])
    response = ""
    for ev in runner.run(user_id="user", session_id=session_id, new_message=content):
        if ev.is_final_response():
            response = ev.content.parts[0].text
    return response

# === Gradio UI ===
with gr.Blocks(title="경희대 미디어학과 AI 상담사", theme="soft") as app:
    gr.Markdown("# 🎬 경희대 미디어학과 AI 상담사")
    chatbot = gr.Chatbot(type="messages", height=400)
    msg = gr.Textbox(show_label=False, placeholder="이 곳에 질문을 입력하세요...")
    send = gr.Button("전송")

    def chat_fn(user_input, history):
        history = history or []
        # 입력 전처리: 오용 방지
        user_input = user_input.replace("전공", "분야").replace("교수", "교수진")
        # RAG 컨텍스트
        docs = rag_system.search_similar_docs(user_input)
        # Combine existing history (dicts) with new user message
        new_history = history + [{"role": "user", "content": user_input}]
        # Insert docs as system messages
        new_history += docs
        # Build prompt text from history
        prompt = "\n".join([f"{m['role']}: {m['content']}" for m in new_history])
        resp = asyncio.run(get_response(prompt))
        new_history.append({"role": "assistant", "content": resp})
        return new_history, ""

    send.click(chat_fn, inputs=[msg, chatbot], outputs=[chatbot, msg])
    msg.submit(chat_fn, inputs=[msg, chatbot], outputs=[chatbot, msg])

    gr.Markdown(f"""
    ---
    ### ⚙️ 시스템 정보\n
    **ChromaDB 문서 수**: {rag_system.collection.count()}개\n
    **임베딩 모델**: snunlp/KR-SBERT-V40K-klueNLI-augSTS (한국어 특화)\n    
    **언어 모델**: Google Gemini 2.0 Flash (무료)  
    """)

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT",7860)), share=False, show_api=False)