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

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

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)},
        )
        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",
        )
        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))
        try:
            r = subprocess.run(
                cmd, shell=True, capture_output=True, text=True,
                timeout=timeout, executable="/bin/bash",
            )
            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", ""))
        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", ""))
        _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)