import os
import chromadb
from sentence_transformers import SentenceTransformer
from google import genai
import gradio as gr

# === 환경 설정 ===
DB_DIR = os.getenv("CHROMA_DB_DIR", os.path.join(os.getcwd(), "chromadb_KH_media"))
os.environ["CHROMA_DB_DIR"] = DB_DIR
API_KEY = os.getenv("GOOGLE_API_KEY", "AIzaSyCoglAa_T_27Qu-nVULgvsV9oPlJxNGS2k")

# === Simple RAG 시스템 ===
class SimpleRAGSystem:
    def __init__(self, db_path=None, collection_name="KH_media_docs"):
        path = db_path or DB_DIR
        self.encoder = SentenceTransformer("snunlp/KR-SBERT-V40K-klueNLI-augSTS")
        self.client = chromadb.PersistentClient(path=path)
        self.collection = self.client.get_collection(name=collection_name)
        self.available = self.collection.count() > 0

    def search(self, query, top_k=10):
        if not self.available:
            return []
        emb = self.encoder.encode(query).tolist()
        result = self.collection.query(
            query_embeddings=[emb],
            n_results=top_k,
            include=["documents"]
        )
        return result.get("documents", [[]])[0]

rag = SimpleRAGSystem()

# === Google GenAI 클라이언트 ===
client = genai.Client(api_key=API_KEY)

# === 시스템 메시지 ===
SYSTEM_MSG = """
당신은 경희대학교 미디어학과 전문 상담 AI입니다.
# 주요 역할:
- 제공된 문서 정보를 바탕으로 답변 제공
- 미디어학과 관련 질문에 친절하고 구체적으로 응답
- 문서에 없는 내용은 일반 지식으로 보완 (단, 명시)
# 답변 스타일:
- 자세하고 풍부한 설명을 포함하여 상세하고 길게 답변 제공
- 친근하고 도움이 되는 상담사 톤
- 핵심 정보를 명확하게 전달
- 추가 궁금한 점이 있으면 언제든 물어보라고 안내
# 참고 문서 활용:
- 문서 내용이 있으면 구체적으로 인용
- 여러 문서의 정보를 종합하여 답변 작성
- 정확하지 않은 정보는 추측하지 말고 솔직하게 모른다고 답변
# 현재 경희대학교 미디어학과 교수진:
이인희, 김태용, 박종민, 홍지아, 이정교, 이기형, 이선영, 조수영, 이종혁, 이두황, 이상원, 이훈, 최수진, 최민아, 김관호
"""

# === 응답 함수 ===
def respond(message, history, system_message, max_tokens, temperature, top_p, model_name):
    docs = rag.search(message) if rag.available else []
    ctx = "\n".join(f"참고문서{i+1}: {d}" for i, d in enumerate(docs))
    sys_msg = system_message + ("\n# 참고문서:\n" + ctx if ctx else "")
    convo = "".join(f"사용자: {u}\nAI: {a}\n" for u, a in history)
    prompt = f"{sys_msg}\n{convo}사용자: {message}\nAI:"
    try:
        response = client.models.generate_content(
            model=model_name,
            contents=prompt,
            config={
                "max_output_tokens": max_tokens,
                "temperature": temperature,
                "top_p": top_p
            }
        )
        return response.text or "응답이 없습니다."
    except Exception as e:
        err = str(e).lower()
        if "quota" in err:
            return "API 할당량을 초과했습니다. 나중에 시도해주세요."
        if "authentication" in err:
            return "인증 오류: API 키를 확인하세요."
        return f"오류 발생: {e}"

# === Gradio 인터페이스 ===
demo = gr.ChatInterface(
    fn=respond,
    title="🎬 경희대학교 미디어학과 AI 상담사",
    description="경희대학교 미디어학과에 대해 물어보세요!",
    additional_inputs=[
        gr.Slider(128, 2048, value=1024, step=64, label="최대 토큰"),
        gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p"),
        gr.Dropdown(
            choices=[
                "gemini-2.0-flash", "gemini-2.0-flash-lite",
                "gemini-1.5-flash", "gemini-1.5-pro",
                "gemma-3-27b-it", "gemma-3-12b-it", "gemma-3-4b-it"
            ],
            value="gemini-2.0-flash",
            label="모델 선택"
        )
    ],
    additional_inputs_accordion="🔧 고급 설정",
    examples=[
        ["미디어학과에서 배우는 주요 과목들은 무엇인가요?"],
        ["미디어학과 교수진을 소개해주세요."],
        ["미디어학과 졸업 후 진로는 어떻게 되나요?"],
        ["미디어학과 입학 전형에 대해 알려주세요."],
        ["미디어학과 동아리나 학생 활동은 어떤 것들이 있나요?"]
    ],
    type="messages",
    theme="soft",
    analytics_enabled=False
)

# === 시스템 정보 표시 ===
with demo:
    with gr.Row():
        with gr.Column():
            gr.Markdown(f"""
            ---
            ### ⚙️ 시스템 정보
            **언어 모델**: Google Gemini 2.0 Flash, Gemma 3 (4B/12B/27B) 선택 가능
            **임베딩 모델**: snunlp/KR-SBERT-V40K-klueNLI-augSTS (한국어 특화)
            **RAG 상태**: {"✅ 활성화" if rag.available else "❌ 비활성화"}
            **문서 수**: {rag.collection.count() if rag.available else "0"}개
            """)

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