# mcp_server.py
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
import os
import json
import asyncio
from typing import Dict, Any

app = FastAPI(title="Simple MCP Server - Example")

# --- Metadata endpoint (tools discovery) ---
@app.get("/.well-known/mcp")
async def mcp_discovery():
    """
    MCP discovery endpoint: لیست ابزارها و metadata را برمی‌گرداند.
    (فرمت دقیق schema بسته به spec ممکن است فرق کند — این یک مثال عملی است.)
    """
    tools = {
        "tools": [
            {
                "name": "list_files",
                "description": "List files in a directory",
                "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}
            },
            {
                "name": "read_file",
                "description": "Read a file's content (bytes limited)",
                "input_schema": {"type": "object", "properties": {"path": {"type": "string"}, "max_bytes": {"type": "integer"}}}
            }
        ]
    }
    return JSONResponse(content=tools)


# --- Helper: implement tools ---
async def handle_list_files(params: Dict[str, Any]) -> Dict[str, Any]:
    path = params.get("path", ".")
    if not os.path.isdir(path):
        raise HTTPException(status_code=400, detail="path is not a directory")
    items = os.listdir(path)
    return {"files": items}

async def handle_read_file(params: Dict[str, Any]) -> Dict[str, Any]:
    path = params.get("path")
    if not path or not os.path.isfile(path):
        raise HTTPException(status_code=400, detail="file not found")
    max_bytes = int(params.get("max_bytes", 2000))
    with open(path, "rb") as f:
        data = f.read(max_bytes)
    try:
        text = data.decode("utf-8", errors="replace")
    except Exception:
        text = "<binary>"
    return {"content": text}


# --- SSE endpoint to handle requests/streams (simple) ---
@app.post("/invoke")
async def invoke_tool(request: Request):
    """
    یک endpoint ساده HTTP که درخواست‌های JSON حاوی:
    { "tool": "<tool_name>", "params": { ... } }
    را قبول می‌کند و پاسخ JSON برمی‌گرداند.
    (برای SSE یا streamable HTTP می‌توانیم از مسیر متفاوت استفاده کنیم.)
    """
    payload = await request.json()
    tool = payload.get("tool")
    params = payload.get("params", {})

    if tool == "list_files":
        result = await handle_list_files(params)
    elif tool == "read_file":
        result = await handle_read_file(params)
    else:
        raise HTTPException(status_code=400, detail="unknown tool")

    return JSONResponse(content={"result": result})


# Optional: a stream endpoint (SSE style) — minimal example
@app.get("/stream")
async def stream_endpoint():
    """
    نمونهٔ بسیار ساده SSE که می‌تواند برای ارتباط streamable طبق MCP استفاده شود.
    برای production باید header های مربوط به SSE و keep-alive مدیریت شود.
    """
    async def event_generator():
        for i in range(3):
            await asyncio.sleep(1)
            yield f"data: {json.dumps({'tick': i})}\n\n"
    return StreamingResponse(event_generator(), media_type="text/event-stream")