import asyncio
import logging
import uvicorn
import signal
import os
import sys
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)

    # Capture the host header to ensure image links work across the network
    host_header = request.headers.get("host", f"{config.HOST}:{config.PORT}")
    config.request_host.set(host_header)

    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__":
    # Use a separate variable for uvicorn config to avoid shadowing the config module
    uvicorn_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(uvicorn_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:
        logger.info("Server process exiting.")
        sys.exit(0)