import gradio as gr
import requests
import json

MCP_URL = "https://huggingface.co/mcp"
HEADERS = {
    "Content-Type": "application/json",
    "Accept": "application/json, text/event-stream"
}

def call_mcp_tool(method, params):
    payload = {
        "jsonrpc": "2.0",
        "method": method,
        "params": params,
        "id": 1
    }
    response = requests.post(MCP_URL, headers=HEADERS, data=json.dumps(payload))
    try:
        return response.json()
    except Exception as e:
        return {"error": str(e), "raw_response": response.text}

def space_search(query, limit=10, mcp=False):
    params = {
        "query": query,
        "limit": limit,
        "mcp": mcp
    }
    result = call_mcp_tool("tools/call", {"tool": "space_search", "input": params})
    return result

def model_detail(model_id):
    params = {
        "model_id": model_id
    }
    result = call_mcp_tool("tools/call", {"tool": "model_detail", "input": params})
    return result

def paper_search(query, results_limit=12, concise_only=False):
    params = {
        "query": query,
        "results_limit": results_limit,
        "concise_only": concise_only
    }
    result = call_mcp_tool("tools/call", {"tool": "paper_search", "input": params})
    return result

def dataset_search(query="", author="", tags=None, limit=20, sort="downloads"):
    if tags is None:
        tags = []
    params = {
        "query": query,
        "author": author,
        "tags": tags,
        "limit": limit,
        "sort": sort
    }
    result = call_mcp_tool("tools/call", {"tool": "dataset_search", "input": params})
    return result

def dataset_detail(dataset_id):
    params = {
        "dataset_id": dataset_id
    }
    result = call_mcp_tool("tools/call", {"tool": "dataset_detail", "input": params})
    return result

def duplicate_space(sourceSpaceId, newSpaceId="", hardware="freecpu", private=True):
    params = {
        "sourceSpaceId": sourceSpaceId,
        "newSpaceId": newSpaceId,
        "hardware": hardware,
        "private": private
    }
    result = call_mcp_tool("tools/call", {"tool": "duplicate_space", "input": params})
    return result

with gr.Blocks() as demo:
    gr.Markdown("# MCP Gradio Client")

    with gr.Tab("Space Search"):
        query = gr.Textbox(label="Query")
        limit = gr.Slider(1, 50, value=10, label="Limit")
        mcp = gr.Checkbox(label="MCP Only")
        output = gr.JSON()
        btn = gr.Button("Search")
        btn.click(space_search, inputs=[query, limit, mcp], outputs=output)

    with gr.Tab("Model Detail"):
        model_id = gr.Textbox(label="Model ID")
        output = gr.JSON()
        btn = gr.Button("Get Details")
        btn.click(model_detail, inputs=model_id, outputs=output)

    with gr.Tab("Paper Search"):
        query = gr.Textbox(label="Query")
        results_limit = gr.Slider(1, 50, value=12, label="Results Limit")
        concise_only = gr.Checkbox(label="Concise Only")
        output = gr.JSON()
        btn = gr.Button("Search")
        btn.click(paper_search, inputs=[query, results_limit, concise_only], outputs=output)

    with gr.Tab("Dataset Search"):
        query = gr.Textbox(label="Query")
        author = gr.Textbox(label="Author")
        tags = gr.Textbox(label="Tags (comma-separated)")
        limit = gr.Slider(1, 100, value=20, label="Limit")
        sort = gr.Dropdown(choices=["downloads", "likes", "lastModified"], value="downloads", label="Sort By")
        output = gr.JSON()
        btn = gr.Button("Search")
        def parse_tags(tags_str):
            return [tag.strip() for tag in tags_str.split(",") if tag.strip()]
        btn.click(lambda q, a, t, l, s: dataset_search(q, a, parse_tags(t), l, s),
                  inputs=[query, author, tags, limit, sort], outputs=output)

    with gr.Tab("Dataset Detail"):
        dataset_id = gr.Textbox(label="Dataset ID")
        output = gr.JSON()
        btn = gr.Button("Get Details")
        btn.click(dataset_detail, inputs=dataset_id, outputs=output)

    with gr.Tab("Duplicate Space"):
        sourceSpaceId = gr.Textbox(label="Source Space ID")
        newSpaceId = gr.Textbox(label="New Space ID")
        hardware = gr.Dropdown(choices=["freecpu", "zerogpu"], value="freecpu", label="Hardware")
        private = gr.Checkbox(label="Private", value=True)
        output = gr.JSON()
        btn = gr.Button("Duplicate")
        btn.click(duplicate_space, inputs=[sourceSpaceId, newSpaceId, hardware, private], outputs=output)

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