"""
MCP HTTP Server — persistent PTY sessions + background jobs
"""

import asyncio
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, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel

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

# ─── 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 = ""
                for x in range(self.cols):
                    char = self.screen.buffer[y][x]
                    row += char.data
                lines.append(row.rstrip())
            # trim trailing blank lines
            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] = {}

# ─── MCP Tool Discovery ───────────────────────────────────────────────────────

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": {}},
    },
]

# ─── MCP Endpoints ────────────────────────────────────────────────────────────

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


@app.get("/mcp/tools")
def list_tools():
    return {"tools": TOOLS}


class ToolCallRequest(BaseModel):
    name: str
    arguments: dict[str, Any] = {}


@app.post("/mcp/call")
async def call_tool(req: ToolCallRequest):
    name = req.name
    args = req.arguments

    # ── exec ─────────────────────────────────────────────────────────────────
    if name == "exec":
        cmd = args.get("cmd", "")
        timeout = int(args.get("timeout", 30))
        try:
            result = subprocess.run(
                cmd, shell=True, capture_output=True, text=True,
                timeout=timeout, executable="/bin/bash",
            )
            return _ok({"stdout": result.stdout, "stderr": result.stderr, "exit_code": result.returncode})
        except subprocess.TimeoutExpired:
            return _err(f"Command timed out after {timeout}s")

    # ── pty_create ────────────────────────────────────────────────────────────
    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)  # let bash prompt appear
        return _ok({"session_id": sess.id, "cols": cols, "rows": rows})

    # ── pty_send ──────────────────────────────────────────────────────────────
    elif name == "pty_send":
        sess = _get_pty(args.get("session_id", ""))
        text = _decode_escapes(args.get("text", ""))
        sess.send(text)
        return _ok({"sent": len(text)})

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

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

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

    # ── job_run ───────────────────────────────────────────────────────────────
    elif name == "job_run":
        job = BackgroundJob(args.get("cmd", ""))
        _jobs[job.id] = job
        return _ok({"job_id": job.id, "status": job.status})

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

    # ── job_status ────────────────────────────────────────────────────────────
    elif name == "job_status":
        job = _get_job(args.get("job_id", ""))
        _, total = job.output(0)
        return _ok({
            "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,
        })

    # ── job_kill ──────────────────────────────────────────────────────────────
    elif name == "job_kill":
        job = _get_job(args.get("job_id", ""))
        job.kill()
        return _ok({"killed": job.id})

    # ── job_list ──────────────────────────────────────────────────────────────
    elif name == "job_list":
        return _ok({
            "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 HTTPException(status_code=404, detail=f"Unknown tool: {name}")


# ─── 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:
        s = m.group(0)
        return s.encode("utf-8").decode("unicode_escape")
    return _ESCAPE_RE.sub(_replace, text)


def _ok(data: dict) -> JSONResponse:
    return JSONResponse({"ok": True, **data})


def _err(msg: str, status: int = 400) -> JSONResponse:
    return JSONResponse({"ok": False, "error": msg}, status_code=status)


def _get_pty(session_id: str) -> PtySession:
    sess = _pty_sessions.get(session_id)
    if not sess:
        raise HTTPException(status_code=404, detail=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 HTTPException(status_code=404, detail=f"Job not found: {job_id}")
    return job


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