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

import asyncio
import fnmatch
import json
import logging
import os
import pathlib
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 = "/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 inside /workspace and return stdout+stderr synchronously (max 30 s). Use cwd to run in a subdirectory.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "cmd": {"type": "string", "description": "Shell command to execute"},
                "cwd": {"type": "string", "description": "Working directory relative to /workspace (default: /workspace)"},
                "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": {}},
    },
    {
        "name": "file_read",
        "description": "Read a file inside /workspace. Returns content with 1-based line numbers (cat -n format). Use offset and limit to read large files in chunks.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "Path relative to /workspace"},
                "offset": {"type": "integer", "description": "Line number to start reading from (1-based, default: 1)"},
                "limit": {"type": "integer", "description": "Maximum number of lines to return"},
            },
            "required": ["path"],
        },
    },
    {
        "name": "file_write",
        "description": "Write (create or overwrite) a file inside /workspace with the given content. Prefer file_edit for modifying existing files.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "Path relative to /workspace"},
                "content": {"type": "string", "description": "Full file content to write"},
            },
            "required": ["path", "content"],
        },
    },
    {
        "name": "file_edit",
        "description": "Perform an exact string replacement in a file inside /workspace. old_string must appear exactly once unless replace_all is true. Fails if old_string is not found or is ambiguous.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "Path relative to /workspace"},
                "old_string": {"type": "string", "description": "Exact text to find and replace"},
                "new_string": {"type": "string", "description": "Replacement text"},
                "replace_all": {"type": "boolean", "description": "Replace all occurrences (default: false)", "default": False},
            },
            "required": ["path", "old_string", "new_string"],
        },
    },
    {
        "name": "file_list",
        "description": "List files and directories inside /workspace (optionally a subdirectory).",
        "inputSchema": {
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "Subdirectory relative to /workspace (default: root)", "default": ""},
            },
        },
    },
    {
        "name": "file_delete",
        "description": "Delete a file inside /workspace.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "Path relative to /workspace"},
            },
            "required": ["path"],
        },
    },
    {
        "name": "file_move",
        "description": "Move or rename a file or directory inside /workspace.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "src": {"type": "string", "description": "Source path relative to /workspace"},
                "dst": {"type": "string", "description": "Destination path relative to /workspace"},
            },
            "required": ["src", "dst"],
        },
    },
    {
        "name": "file_glob",
        "description": "Find files inside /workspace matching a glob pattern (e.g. '**/*.py'). Returns matching paths relative to /workspace.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "pattern": {"type": "string", "description": "Glob pattern (e.g. '**/*.py', 'src/**/*.ts')"},
                "path": {"type": "string", "description": "Subdirectory to search in (default: workspace root)", "default": ""},
            },
            "required": ["pattern"],
        },
    },
    {
        "name": "file_grep",
        "description": "Search file contents inside /workspace using a regex pattern. Returns matching lines with file path and line number.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "pattern": {"type": "string", "description": "Regular expression to search for"},
                "path": {"type": "string", "description": "File or directory to search in (default: workspace root)", "default": ""},
                "glob": {"type": "string", "description": "Only search files matching this glob (e.g. '*.py')"},
                "context": {"type": "integer", "description": "Lines of context before and after each match (default: 0)", "default": 0},
                "ignore_case": {"type": "boolean", "description": "Case-insensitive search (default: false)", "default": False},
                "max_matches": {"type": "integer", "description": "Maximum number of matches to return (default: 100)", "default": 100},
            },
            "required": ["pattern"],
        },
    },
]

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

async def _dispatch(name: str, args: dict) -> dict:
    if name == "exec":
        cmd = args.get("cmd", "")
        timeout = int(args.get("timeout", 30))
        cwd = _workspace_path(args["cwd"]) if args.get("cwd") else pathlib.Path(WORKSPACE)
        log.info("exec (cwd=%s): %s", cwd, cmd)
        try:
            r = subprocess.run(
                cmd, shell=True, capture_output=True, text=True,
                timeout=timeout, executable="/bin/bash",
                cwd=str(cwd),
            )
            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()
            ]
        }

    elif name == "file_read":
        p = _workspace_path(args["path"])
        log.info("file_read: %s", p)
        all_lines = p.read_text(encoding="utf-8").splitlines()
        offset = max(1, int(args.get("offset", 1))) - 1  # convert to 0-based
        limit = args.get("limit")
        sliced = all_lines[offset:offset + int(limit)] if limit is not None else all_lines[offset:]
        numbered = "\n".join(f"{offset + i + 1}\t{line}" for i, line in enumerate(sliced))
        return {"path": args["path"], "content": numbered, "total_lines": len(all_lines)}

    elif name == "file_write":
        p = _workspace_path(args["path"])
        log.info("file_write: %s (%d bytes)", p, len(args.get("content", "")))
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(args.get("content", ""), encoding="utf-8")
        return {"path": args["path"], "written": len(args.get("content", ""))}

    elif name == "file_edit":
        p = _workspace_path(args["path"])
        old_string = args["old_string"]
        new_string = args["new_string"]
        replace_all = bool(args.get("replace_all", False))
        log.info("file_edit: %s replace_all=%s", p, replace_all)
        content = p.read_text(encoding="utf-8")
        count = content.count(old_string)
        if count == 0:
            raise ValueError(f"old_string not found in {args['path']}")
        if not replace_all and count > 1:
            raise ValueError(
                f"old_string is not unique in {args['path']} ({count} occurrences). "
                "Use replace_all=true to replace all, or provide more context to make it unique."
            )
        if replace_all:
            new_content = content.replace(old_string, new_string)
        else:
            new_content = content.replace(old_string, new_string, 1)
        p.write_text(new_content, encoding="utf-8")
        return {"path": args["path"], "replacements": count if replace_all else 1}

    elif name == "file_list":
        p = _workspace_path(args.get("path", ""))
        log.info("file_list: %s", p)
        entries = [
            {"name": e.name, "type": "dir" if e.is_dir() else "file", "size": e.stat().st_size if e.is_file() else None}
            for e in sorted(p.iterdir(), key=lambda e: (e.is_file(), e.name))
        ]
        return {"path": args.get("path", ""), "entries": entries}

    elif name == "file_delete":
        p = _workspace_path(args["path"])
        log.info("file_delete: %s", p)
        p.unlink()
        return {"deleted": args["path"]}

    elif name == "file_move":
        src = _workspace_path(args["src"])
        dst = _workspace_path(args["dst"])
        log.info("file_move: %s -> %s", src, dst)
        dst.parent.mkdir(parents=True, exist_ok=True)
        src.rename(dst)
        return {"src": args["src"], "dst": args["dst"]}

    elif name == "file_glob":
        base = _workspace_path(args.get("path", ""))
        pattern = args["pattern"]
        log.info("file_glob: %s in %s", pattern, base)
        matches = sorted(
            str(p.relative_to(pathlib.Path(WORKSPACE)))
            for p in base.glob(pattern)
            if p.is_file()
        )
        return {"pattern": pattern, "matches": matches, "count": len(matches)}

    elif name == "file_grep":
        base = _workspace_path(args.get("path", ""))
        pattern = args["pattern"]
        glob_filter = args.get("glob")
        context_lines = int(args.get("context", 0))
        ignore_case = bool(args.get("ignore_case", False))
        max_matches = int(args.get("max_matches", 100))
        log.info("file_grep: %r in %s", pattern, base)
        flags = re.IGNORECASE if ignore_case else 0
        regex = re.compile(pattern, flags)
        results = []
        files = base.rglob("*") if base.is_dir() else [base]
        for filepath in sorted(files):
            if not filepath.is_file():
                continue
            if glob_filter and not fnmatch.fnmatch(filepath.name, glob_filter):
                continue
            try:
                lines = filepath.read_text(encoding="utf-8", errors="replace").splitlines()
            except OSError:
                continue
            for i, line in enumerate(lines):
                if regex.search(line):
                    rel = str(filepath.relative_to(pathlib.Path(WORKSPACE)))
                    entry = {"file": rel, "line": i + 1, "match": line}
                    if context_lines:
                        entry["before"] = lines[max(0, i - context_lines):i]
                        entry["after"] = lines[i + 1:i + 1 + context_lines]
                    results.append(entry)
                    if len(results) >= max_matches:
                        return {"pattern": pattern, "matches": results, "truncated": True}
        return {"pattern": pattern, "matches": results, "truncated": False}

    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 _workspace_path(rel: str) -> pathlib.Path:
    p = pathlib.Path(rel)
    return p if p.is_absolute() else pathlib.Path(WORKSPACE) / rel


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, access_log=False)