from fastmcp import FastMCP
from typing import Literal

# FastMCP 서버 초기화
mcp = FastMCP(
    "RAG_Server",
    instructions="A RAG server that provides vector search, document addition, and web search capabilities."
)

# 간단한 메모리 기반 벡터 스토어
vector_store = []

@mcp.tool()
async def vector_search(
    query: str,
    search_type: Literal["semantic", "keyword", "hybrid"] = "semantic",
    k: int = 5
) -> str:
    """벡터 스토어에서 문서를 검색합니다."""
    global vector_store

    if not vector_store:
        return "벡터 스토어가 비어있습니다. 먼저 add_document를 사용하여 문서를 추가하세요."

    # 간단한 키워드 검색
    if search_type == "keyword":
        results = [doc for doc in vector_store if query.lower() in doc.lower()][:k]
    else:
        # 시맨틱 검색 시뮬레이션
        results = vector_store[:k]

    return "\n\n".join(results[:k]) if results else "검색 결과가 없습니다."

@mcp.tool()
async def add_document(text: str, metadata: dict = None) -> str:
    """사용자 텍스트를 벡터 스토어에 추가합니다."""
    global vector_store

    if metadata is None:
        metadata = {"source": "user_input"}

    vector_store.append(text)

    return f"문서가 성공적으로 추가되었습니다. 총 {len(text)} 문자, 현재 저장된 문서 수: {len(vector_store)}"

@mcp.tool()
async def web_search(query: str, max_results: int = 3) -> str:
    """웹 검색을 시뮬레이션합니다."""
    return f"웹 검색 시뮬레이션: '{query}'에 대한 검색 결과입니다. (실제 웹 검색은 API 키가 필요합니다)"

def main():
    """Main entry point for the MCP server."""
    print("RAG MCP 서버를 초기화합니다...")
    print("벡터 스토어 초기화 완료!")
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()