#!/usr/bin/env python3
"""
HTTP/SSE Transport Wrapper for SSH-MCP Server

This provides an HTTP interface with Server-Sent Events (SSE) support
for Claude Code to connect to the SSH-MCP server.
"""

import asyncio
import json
import subprocess
from typing import Optional
from fastapi import FastAPI, Request, Response, HTTPException
from fastapi.responses import StreamingResponse
import uvicorn

app = FastAPI(title="SSH-MCP Server", version="1.0.0")

# In-memory storage for MCP server process
mcp_process: Optional[asyncio.subprocess.Process] = None


async def start_mcp_server():
    """Start the MCP server subprocess"""
    global mcp_process
    if mcp_process is None:
        mcp_process = await asyncio.create_subprocess_exec(
            "/usr/bin/python3",
            "/opt/ssh-mcp-server/server.py",
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
    return mcp_process


async def send_mcp_request(request_data: dict) -> dict:
    """Send request to MCP server and get response"""
    process = await start_mcp_server()

    # Send request
    request_json = json.dumps(request_data) + "\n"
    process.stdin.write(request_json.encode())
    await process.stdin.drain()

    # Read response
    response_line = await process.stdout.readline()
    response_data = json.loads(response_line.decode())

    return response_data


@app.get("/")
async def root():
    """Health check endpoint"""
    return {
        "name": "SSH-MCP Server",
        "version": "1.0.0",
        "status": "operational",
        "description": "MCP server providing SSH-based infrastructure access",
        "endpoints": {
            "/": "Health check",
            "/tools": "List available tools",
            "/tools/call": "Call a tool",
            "/sse": "Server-Sent Events endpoint"
        }
    }


@app.get("/health")
async def health():
    """Health check"""
    return {"status": "healthy"}


@app.post("/tools")
async def list_tools():
    """List available MCP tools"""
    request = {
        "jsonrpc": "2.0",
        "id": 1,
        "method": "tools/list",
        "params": {}
    }

    response = await send_mcp_request(request)
    return response.get("result", {})


@app.post("/tools/call")
async def call_tool(request: Request):
    """Call an MCP tool"""
    body = await request.json()

    tool_name = body.get("name")
    arguments = body.get("arguments", {})

    mcp_request = {
        "jsonrpc": "2.0",
        "id": 2,
        "method": "tools/call",
        "params": {
            "name": tool_name,
            "arguments": arguments
        }
    }

    response = await send_mcp_request(mcp_request)
    return response.get("result", {})


@app.get("/sse")
async def sse_endpoint(request: Request):
    """
    Server-Sent Events endpoint for real-time MCP communication
    This is required for Claude Code's MCP connector
    """
    async def event_generator():
        # Send initial connection event
        yield f"data: {json.dumps({'type': 'connected', 'server': 'ssh-mcp-server'})}\n\n"

        # Keep connection alive
        while True:
            if await request.is_disconnected():
                break

            # Send periodic heartbeat
            yield f"data: {json.dumps({'type': 'heartbeat', 'timestamp': asyncio.get_event_loop().time()})}\n\n"
            await asyncio.sleep(30)

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
        }
    )


@app.post("/message")
async def message(request: Request):
    """
    MCP message endpoint for Claude Code integration
    Accepts JSON-RPC 2.0 messages
    """
    body = await request.json()
    response = await send_mcp_request(body)
    return response


# MCP Server Info endpoint (required by MCP spec)
@app.get("/mcp/info")
async def mcp_info():
    """MCP server information"""
    return {
        "name": "ssh-mcp-server",
        "version": "1.0.0",
        "description": "SSH-based MCP server for OBNH/OBR infrastructure management",
        "capabilities": {
            "tools": True,
            "resources": False,
            "prompts": False
        },
        "transport": "http+sse"
    }


if __name__ == "__main__":
    print("Starting SSH-MCP HTTP/SSE Server on port 8081...")
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8081,
        log_level="info"
    )