"""
MCP HTTP Server — persistent PTY sessions + background jobs
Implements MCP Streamable HTTP transport (JSON-RPC 2.0)
"""

import asyncio
import json
import logging
import os
import re
import signal
import subprocess
import threading
import time
import uuid
from typing import Any

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
log = logging.getLogger("mcp")

WORKSPACE = os.environ.get("WORKSPACE_DIR", "/workspace")

import pyte
import ptyprocess
import uvicorn
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

app = FastAPI(title="MCP Bash Server", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ─── Stores ──────────────────────────────────────────────────────────────────

class PtySession:
    def __init__(self, cols: int = 220, rows: int = 50):
        self.id = str(uuid.uuid4())
        self.cols = cols
        self.rows = rows
        self.screen = pyte.Screen(cols, rows)
        self.stream = pyte.ByteStream(self.screen)
        self._lock = threading.Lock()
        self.created_at = time.time()

        self.proc = ptyprocess.PtyProcess.spawn(
            ["/bin/bash", "--login"],
            dimensions=(rows, cols),
            env={**os.environ, "TERM": "xterm-256color", "COLUMNS": str(cols), "LINES": str(rows)},
            cwd=WORKSPACE,
        )
        self._reader = threading.Thread(target=self._read_loop, daemon=True)
        self._reader.start()

    def _read_loop(self):
        while self.proc.isalive():
            try:
                data = self.proc.read(4096)
            except EOFError:
                break
            with self._lock:
                self.stream.feed(data)

    def send(self, text: str):
        self.proc.write(text.encode())

    def snapshot(self) -> dict[str, Any]:
        with self._lock:
            lines = []
            for y in range(self.rows):
                row = "".join(self.screen.buffer[y][x].data for x in range(self.cols))
                lines.append(row.rstrip())
            while lines and lines[-1] == "":
                lines.pop()
            cursor = {"x": self.screen.cursor.x, "y": self.screen.cursor.y}
        return {"lines": lines, "cursor": cursor, "alive": self.proc.isalive()}

    def kill(self):
        try:
            self.proc.terminate(force=True)
        except Exception:
            pass


class BackgroundJob:
    def __init__(self, cmd: str):
        self.id = str(uuid.uuid4())
        self.cmd = cmd
        self.output_lines: list[str] = []
        self.exit_code: int | None = None
        self.started_at = time.time()
        self.finished_at: float | None = None
        self._lock = threading.Lock()

        self._proc = subprocess.Popen(
            cmd,
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            executable="/bin/bash",
            cwd=WORKSPACE,
        )
        self._reader = threading.Thread(target=self._read_loop, daemon=True)
        self._reader.start()

    def _read_loop(self):
        for line in self._proc.stdout:
            with self._lock:
                self.output_lines.append(line.rstrip("\n"))
        self._proc.wait()
        with self._lock:
            self.exit_code = self._proc.returncode
            self.finished_at = time.time()

    @property
    def status(self) -> str:
        if self.exit_code is None:
            return "running"
        return "done" if self.exit_code == 0 else "failed"

    def output(self, offset: int = 0) -> tuple[list[str], int]:
        with self._lock:
            return self.output_lines[offset:], len(self.output_lines)

    def kill(self):
        try:
            self._proc.send_signal(signal.SIGTERM)
        except Exception:
            pass


_pty_sessions: dict[str, PtySession] = {}
_jobs: dict[str, BackgroundJob] = {}

# ─── Tool definitions ─────────────────────────────────────────────────────────

TOOLS = [
    {
        "name": "exec",
        "description": "Run a shell command and return stdout+stderr synchronously (max 30 s).",
        "inputSchema": {
            "type": "object",
            "properties": {
                "cmd": {"type": "string", "description": "Shell command to execute"},
                "timeout": {"type": "integer", "description": "Timeout in seconds (default 30)", "default": 30},
            },
            "required": ["cmd"],
        },
    },
    {
        "name": "pty_create",
        "description": "Open a persistent PTY (interactive terminal) session. Returns session_id.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "cols": {"type": "integer", "default": 220},
                "rows": {"type": "integer", "default": 50},
            },
        },
    },
    {
        "name": "pty_send",
        "description": "Send text/keystrokes to a PTY session. Use \\n for Enter, \\x03 for Ctrl-C, \\x04 for Ctrl-D.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "session_id": {"type": "string"},
                "text": {"type": "string", "description": "Raw text to send (escape sequences allowed)"},
            },
            "required": ["session_id", "text"],
        },
    },
    {
        "name": "pty_read",
        "description": "Read the current screen state of a PTY session (rendered terminal grid).",
        "inputSchema": {
            "type": "object",
            "properties": {
                "session_id": {"type": "string"},
                "wait_ms": {"type": "integer", "default": 200, "description": "Wait N ms before reading (let output settle)"},
            },
            "required": ["session_id"],
        },
    },
    {
        "name": "pty_kill",
        "description": "Terminate and remove a PTY session.",
        "inputSchema": {
            "type": "object",
            "properties": {"session_id": {"type": "string"}},
            "required": ["session_id"],
        },
    },
    {
        "name": "pty_list",
        "description": "List all open PTY sessions.",
        "inputSchema": {"type": "object", "properties": {}},
    },
    {
        "name": "job_run",
        "description": "Start a long-running background job. Returns job_id immediately.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "cmd": {"type": "string", "description": "Shell command to run in background"},
            },
            "required": ["cmd"],
        },
    },
    {
        "name": "job_output",
        "description": "Retrieve output lines of a background job (optionally from an offset).",
        "inputSchema": {
            "type": "object",
            "properties": {
                "job_id": {"type": "string"},
                "offset": {"type": "integer", "default": 0, "description": "Start reading from this line index"},
            },
            "required": ["job_id"],
        },
    },
    {
        "name": "job_status",
        "description": "Get status (running/done/failed), exit code, and line count of a background job.",
        "inputSchema": {
            "type": "object",
            "properties": {"job_id": {"type": "string"}},
            "required": ["job_id"],
        },
    },
    {
        "name": "job_kill",
        "description": "Send SIGTERM to a running background job.",
        "inputSchema": {
            "type": "object",
            "properties": {"job_id": {"type": "string"}},
            "required": ["job_id"],
        },
    },
    {
        "name": "job_list",
        "description": "List all known background jobs with their current status.",
        "inputSchema": {"type": "object", "properties": {}},
    },
]

# ─── Tool dispatch ────────────────────────────────────────────────────────────

async def _dispatch(name: str, args: dict) -> dict:
    if name == "exec":
        cmd = args.get("cmd", "")
        timeout = int(args.get("timeout", 30))
        log.info("exec: %s", cmd)
        try:
            r = subprocess.run(
                cmd, shell=True, capture_output=True, text=True,
                timeout=timeout, executable="/bin/bash",
                cwd=WORKSPACE,
            )
            return {"stdout": r.stdout, "stderr": r.stderr, "exit_code": r.returncode}
        except subprocess.TimeoutExpired:
            return {"error": f"Command timed out after {timeout}s"}

    elif name == "pty_create":
        cols = int(args.get("cols", 220))
        rows = int(args.get("rows", 50))
        sess = PtySession(cols=cols, rows=rows)
        _pty_sessions[sess.id] = sess
        await asyncio.sleep(0.3)
        return {"session_id": sess.id, "cols": cols, "rows": rows}

    elif name == "pty_send":
        sess = _get_pty(args["session_id"])
        text = _decode_escapes(args.get("text", ""))
        log.info("pty_send [%s]: %r", sess.id[:8], text)
        sess.send(text)
        return {"sent": len(text)}

    elif name == "pty_read":
        sess = _get_pty(args["session_id"])
        wait_ms = int(args.get("wait_ms", 200))
        if wait_ms > 0:
            await asyncio.sleep(wait_ms / 1000)
        return sess.snapshot()

    elif name == "pty_kill":
        sess = _get_pty(args["session_id"])
        sess.kill()
        del _pty_sessions[sess.id]
        return {"killed": sess.id}

    elif name == "pty_list":
        return {
            "sessions": [
                {"session_id": sid, "alive": s.proc.isalive(), "created_at": s.created_at}
                for sid, s in _pty_sessions.items()
            ]
        }

    elif name == "job_run":
        job = BackgroundJob(args.get("cmd", ""))
        log.info("job_run [%s]: %s", job.id[:8], job.cmd)
        _jobs[job.id] = job
        return {"job_id": job.id, "status": job.status}

    elif name == "job_output":
        job = _get_job(args["job_id"])
        offset = int(args.get("offset", 0))
        lines, total = job.output(offset)
        return {"job_id": job.id, "offset": offset, "lines": lines, "total_lines": total}

    elif name == "job_status":
        job = _get_job(args["job_id"])
        _, total = job.output(0)
        return {
            "job_id": job.id, "status": job.status,
            "exit_code": job.exit_code, "total_lines": total,
            "started_at": job.started_at, "finished_at": job.finished_at,
        }

    elif name == "job_kill":
        job = _get_job(args["job_id"])
        job.kill()
        return {"killed": job.id}

    elif name == "job_list":
        return {
            "jobs": [
                {"job_id": jid, "cmd": j.cmd, "status": j.status,
                 "exit_code": j.exit_code, "total_lines": j.output(0)[1]}
                for jid, j in _jobs.items()
            ]
        }

    raise ValueError(f"Unknown tool: {name}")


# ─── MCP JSON-RPC 2.0 endpoint ────────────────────────────────────────────────

async def _handle_rpc(body: dict) -> dict | None:
    method = body.get("method", "")
    params = body.get("params") or {}
    req_id = body.get("id")

    try:
        if method == "initialize":
            result = {
                "protocolVersion": "2024-11-05",
                "capabilities": {"tools": {}},
                "serverInfo": {"name": "mcp-bash", "version": "1.0.0"},
            }
        elif method in ("notifications/initialized", "notifications/cancelled"):
            return None  # notifications need no response
        elif method == "ping":
            result = {}
        elif method == "tools/list":
            result = {"tools": TOOLS}
        elif method == "tools/call":
            name = params.get("name", "")
            args = params.get("arguments") or {}
            data = await _dispatch(name, args)
            result = {"content": [{"type": "text", "text": json.dumps(data)}]}
        else:
            return {"jsonrpc": "2.0", "id": req_id,
                    "error": {"code": -32601, "message": f"Method not found: {method}"}}
    except Exception as exc:
        return {"jsonrpc": "2.0", "id": req_id,
                "error": {"code": -32603, "message": str(exc)}}

    if req_id is None:
        return None
    return {"jsonrpc": "2.0", "id": req_id, "result": result}


@app.post("/")
@app.post("/mcp")
async def mcp_endpoint(request: Request):
    body = await request.json()
    response = await _handle_rpc(body)
    if response is None:
        return JSONResponse({}, status_code=202)
    return JSONResponse(response)


@app.get("/")
def health():
    return {"service": "mcp-bash", "version": "1.0.0"}


# ─── Helpers ──────────────────────────────────────────────────────────────────

_ESCAPE_RE = re.compile(r"\\(n|r|t|\\|x[0-9a-fA-F]{2}|u[0-9a-fA-F]{4})")

def _decode_escapes(text: str) -> str:
    """Decode only ASCII escape sequences — leaves UTF-8 text intact."""
    def _replace(m: re.Match) -> str:
        return m.group(0).encode("utf-8").decode("unicode_escape")
    return _ESCAPE_RE.sub(_replace, text)


def _get_pty(session_id: str) -> PtySession:
    sess = _pty_sessions.get(session_id)
    if not sess:
        raise ValueError(f"PTY session not found: {session_id}")
    return sess


def _get_job(job_id: str) -> BackgroundJob:
    job = _jobs.get(job_id)
    if not job:
        raise ValueError(f"Job not found: {job_id}")
    return job


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)