from typing import Any, Dict, List, Optional

import gradio as gr
from core import MCPyLate
from huggingface_hub import snapshot_download

"""
MCPyLate Server
A Model Context Protocol server that provides search functionality using PyLate.
"""

snapshot_download(
    repo_id="lightonai/leetcode_reasonmoderncolbert",
    local_dir="indexes/",
    repo_type="dataset",
)
mcpylate = MCPyLate()


def pylate_search_leetcode(
    query: str, k: int = 10, index_name: Optional[str] = None
) -> List[Dict[str, Any]]:
    """
    Search the PyLate with multi-vector models in the leetcode collection containing code problems solutions and return top-k hits
    Args:
        query: Search query string
        k: Number of results to return (default: 10)
        index_name: Name of index to search (default: use default index)
    Returns:
        List of search results with docid, score, text snippet, and index name
    """
    return mcpylate.search(query, k)


demo = gr.Interface(
    fn=pylate_search_leetcode,
    inputs=["text"],
    outputs="text",
    title="LeetCode Search",
    description="Search in leetcode database index using PyLate",
)

demo.launch(mcp_server=True, share=True)