import asyncio
import logging
import uvicorn
import signal
from starlette.applications import Starlette
from starlette.routing import Route
from starlette.responses import Response, StreamingResponse
from starlette.middleware.cors import CORSMiddleware

from mcp_logic import MCPServer
from tools import TOOL_REGISTRY

# --- Configuration & Logging ---
logging.basicConfig(
    level=logging.DEBUG,
    filename="debug.log",
    filemode="a",
    format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("MattCP")

mcp_logic = MCPServer()

# --- Transport Layer (Starlette/SSE) ---

async def sse_endpoint(request):
    sid = mcp_logic.create_session()
    logger.debug(f"Session started: {sid}")
    
    async def event_generator():
        try:
            yield f"event: endpoint\ndata: /messages?session_id={sid}\n\n"
            queue = mcp_logic.sessions[sid]
            while mcp_logic.running:
                try:
                    msg = await asyncio.wait_for(queue.get(), timeout=1.0)
                    yield f"event: message\ndata: {msg}\n\n"
                    queue.task_done()
                except asyncio.TimeoutError:
                    continue
        except asyncio.CancelledError:
            logger.debug(f"Session cancelled: {sid}")
        finally:
            mcp_logic.sessions.pop(sid, None)
            logger.debug(f"Session closed: {sid}")

    return StreamingResponse(event_generator(), media_type="text/event-stream")

async def messages_endpoint(request):
    sid = request.query_params.get("session_id")
    if not sid or sid not in mcp_logic.sessions:
        return Response("Session not found", status_code=404)

    try:
        body = await request.json()
    except Exception as e:
        return Response(f"Invalid JSON: {e}", status_code=400)

    method = body.get("method")
    params = body.get("params", {})
    req_id = body.get("id")
    
    logger.debug(f"Request: {method} (ID: {req_id})")
    
    result = await mcp_logic.handle_request(method, params, req_id, TOOL_REGISTRY)
    
    if isinstance(result, Response):
        return result
    
    await mcp_logic.send_to_session(sid, {"jsonrpc": "2.0", "id": req_id, "result": result})
    return Response(status_code=202)

app = Starlette(
    routes=[
        Route("/sse", endpoint=sse_endpoint),
        Route("/messages", endpoint=messages_endpoint, methods=["POST"]),
    ]
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

if __name__ == "__main__":
    config = uvicorn.Config(
        app=app, 
        host="127.0.0.1", 
        port=8000, 
        log_level="info",
        timeout_graceful_shutdown=5
    )
    server = uvicorn.Server(config)
    
    def signal_handler(sig, frame):
        logger.info("Shutdown signal received")
        mcp_logic.running = False
    
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    asyncio.run(server.serve())