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

from mcp_logic import MCPServer
from tools import TOOL_REGISTRY
import config


# --- Configuration & Logging ---
logging.basicConfig(
    level=getattr(logging, config.LOG_LEVEL),
    filename=config.LOG_FILE,

    filemode="a",
    format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("MooseCP")

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:
                    # Use a slightly shorter timeout for faster shutdown response
                    msg = await asyncio.wait_for(queue.get(), timeout=0.5)
                    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)

async def file_endpoint(request):
    """
    Proxy endpoint to serve files from disk with CORP/CORS headers to bypass browser restrictions.
    """
    path = request.path_params.get("path")
    if not path:
        return Response("Path not provided", status_code=400)
    
    # Ensure the path is absolute (SDFiles are usually in /tmp/gradio/...)
    # If the path doesn't start with /, we assume it's relative to root for simplicity in this context
    full_path = path if path.startswith("/") else f"/{path}"
    
    if not os.path.exists(full_path):
        return Response(f"File not found: {full_path}", status_code=404)
    
    return FileResponse(
        full_path, 
        headers={
            "Cross-Origin-Resource-Policy": "cross-origin",
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Methods": "GET"
        }
    )

app = Starlette(
    routes=[
        Route("/sse", endpoint=sse_endpoint),
        Route("/messages", endpoint=messages_endpoint, methods=["POST"]),
        Route("/file/{path:path}", endpoint=file_endpoint),
    ]
)

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

if __name__ == "__main__":
    config = uvicorn.Config(
        app=app, 
        host=config.HOST, 
        port=config.PORT, 
        log_level="info",
        timeout_graceful_shutdown=2 # Reduced from 5 to 2 seconds
    )

    server = uvicorn.Server(config)
    
    def signal_handler(sig, frame):
        logger.info("Shutdown signal received")
        mcp_logic.running = False
        # Explicitly tell Uvicorn to stop the server
        server.should_exit = True
    
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    try:
        asyncio.run(server.serve())
    except KeyboardInterrupt:
        pass
    finally:
        # Final fallback to ensure the process actually dies
        # if the event loop is still hanging on a connection
        os._exit(0)