from mcp.server.fastmcp import FastMCP
from tools.law_tool import LawTool
from tools.law_rag_query import LawRAGQuery
import gradio as gr
from datasets import load_dataset

law_tool = LawTool()
law_rag_query = LawRAGQuery()

# Create an MCP server
mcp = FastMCP("Law Tool Service")

dataset = load_dataset("robin0307/law", split='train')

# Tool implementation
@mcp.tool()
def get_law(category: str, number: int) -> str:
    """
    This is a tool that returns law content by input the category and number.

    Args:
        category: the law category (such as 民法, 中華民國刑法, 民事訴訟法, 刑事訴訟法, 律師法 etc).
        number: the law number (such as 23).

    Returns:
        str: The content of the law.
    """
    result = law_tool(category, number)
    return result

# Tool implementation
@mcp.tool()
def search_law(keyword: str) -> list:
    """
    This is a tool that returns law content by input the keyword.

    Args:
        keyword: the keyword would like to search.

    Returns:
        list: The content list of the law.
    """
    results = dataset.filter(lambda x: keyword in x['content'])
    return [[item['category'], item['content']] for item in results]

# Tool implementation
@mcp.tool()
def rag_query(question: str) -> list:
    """
    This is a tool that returns law content by input a question. It will find the related law and return.

    Args:
        question: the question to query the law.

    Returns:
        list: A list of law content related to the question.
    """
    result = law_rag_query(question)
    return result

with gr.Blocks() as demo:
    with gr.Tabs():
        with gr.Tab("Law Tool"):
            with gr.Row():
                category = gr.Dropdown(label="Law Category", choices=["民法", "中華民國刑法", "民事訴訟法", "刑事訴訟法", "律師法"], info="選擇法律類別")
                number = gr.Number(label="Law Number", info="ex：23")
            query_btn = gr.Button("Submit")
            result = gr.Textbox(label="Result")
            query_btn.click(fn=get_law, inputs=[category, number], outputs=result)
        
        with gr.Tab("Law Search"):
            with gr.Row():
                keyword = gr.Textbox(label="Keyword")
            search_btn = gr.Button("Submit")
            search_output = gr.List(headers=["category", "content"], value=[], label="Result", col_count=2)
            search_btn.click(fn=search_law, inputs=keyword, outputs=search_output)

        with gr.Tab("Law RAG Query"):
            with gr.Row():
                text_input = gr.Textbox(label="Question")
            rag_btn = gr.Button("Submit")
            text_output = gr.List(headers=["content", "score"], value=[], label="Result", col_count=2)
            rag_btn.click(fn=rag_query, inputs=text_input, outputs=text_output)

if __name__ == "__main__":
    demo.launch(mcp_server=True,server_name="0.0.0.0",allowed_paths=["/"])