#!/usr/bin/env python3
"""
Turnstone MCP Server — Improved Multi-Tool Server

Provides 4 tools:
  1. execute_python_code  — sandboxed Python execution with timeouts & limits
  2. playwright_navigate  — browser navigation with timeout & validation
  3. playwright_get_page_text — rendered page text extraction with timeout
  4. extract_to_markdown  — format conversion with size limits & validation
"""

import os
import sys
import re
import logging
import asyncio
import subprocess
import tempfile
from typing import Optional
from urllib.parse import urlparse

from fastmcp import FastMCP
from fastmcp.server.middleware import Middleware, MiddlewareContext
from starlette.middleware import Middleware as StarletteMiddleware
from starlette.middleware.cors import CORSMiddleware
from starlette.responses import JSONResponse

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s - %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("turnstone-mcp")

# ---------------------------------------------------------------------------
# Configuration (env-overridable)
# ---------------------------------------------------------------------------
PYTHON_EXEC_TIMEOUT = int(os.environ.get("PYTHON_EXEC_TIMEOUT", "30"))  # seconds
PYTHON_MAX_OUTPUT = int(os.environ.get("PYTHON_MAX_OUTPUT", "65536"))   # bytes
PYTHON_MAX_MEMORY_MB = int(os.environ.get("PYTHON_MAX_MEMORY_MB", "256"))

PLAYWRIGHT_NAV_TIMEOUT = int(os.environ.get("PLAYWRIGHT_NAV_TIMEOUT", "30000"))  # ms
PLAYWRIGHT_TEXT_TIMEOUT = int(os.environ.get("PLAYWRIGHT_TEXT_TIMEOUT", "15000"))  # ms

EXTRACT_MAX_FILE_SIZE = int(os.environ.get("EXTRACT_MAX_FILE_SIZE", "10485760"))  # 10 MB

WORKSPACE_ROOT = os.path.abspath(os.environ.get("WORKSPACE_ROOT", os.getcwd()))

logger.info("Workspace Root: %s", WORKSPACE_ROOT)

# ---------------------------------------------------------------------------
# FastMCP app
# ---------------------------------------------------------------------------
mcp = FastMCP("turnstone-mcp")

# ---------------------------------------------------------------------------
# Playwright lifecycle (stateful browser session)
# ---------------------------------------------------------------------------
class PlaywrightBrowserManager:
    """Encapsulates Playwright state to prevent garbage collection teardowns."""
    
    def __init__(self):
        self.context_manager = None
        self.playwright = None
        self.browser = None
        self.context = None
        self.page = None

    async def get_page(self):
        """Lazily initialize and return a headless Chromium page."""
        if self.page is None:
            from playwright.async_api import async_playwright

            logger.info("Starting headless Playwright session…")
            
            # CRITICAL: We must store the context_manager as a class property.
            # If it's garbage collected, it forces the Playwright driver to close.
            self.context_manager = async_playwright()
            self.playwright = await self.context_manager.start()
            
            self.browser = await self.playwright.chromium.launch(headless=True)
            self.context = await self.browser.new_context(
                viewport={"width": 1280, "height": 720},
                user_agent=(
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                    "AppleWebKit/537.36 (KHTML, like Gecko) "
                    "Chrome/124.0.0.0 Safari/537.36"
                ),
            )
            self.page = await self.context.new_page()
            
        return self.page

# Initialize singleton manager
browser_manager = PlaywrightBrowserManager()

def _resolve_safe_path(relative_path: str) -> str:
    """Resolve a relative path against WORKSPACE_ROOT, blocking traversal."""
    target = os.path.abspath(os.path.join(WORKSPACE_ROOT, relative_path))
    if not target.startswith(WORKSPACE_ROOT):
        raise PermissionError("Access denied: path escapes WORKSPACE_ROOT")
    return target

# ---------------------------------------------------------------------------
# Health endpoint
# ---------------------------------------------------------------------------
@mcp.custom_route("/health", methods=["GET"])
async def health_check(request):
    return JSONResponse({"status": "healthy", "service": "turnstone-mcp"})

# ---------------------------------------------------------------------------
# Semantic-logging middleware
# ---------------------------------------------------------------------------
class ToolObserverMiddleware(Middleware):
    """Log incoming tool invocations and responses."""

    async def on_call_tool(self, context: MiddlewareContext, call_next):
        tool_name = getattr(context.message, "name", "unknown")
        arguments = getattr(context.message, "arguments", {})
        logger.info("📥 [LLM REQUEST] Invoking tool: '%s'", tool_name)
        logger.info("    ↳ Arguments: %s", arguments)
        try:
            result = await call_next(context)
            content = getattr(result, "content", result)
            logger.info("📤 [SERVER RESPONSE] Tool '%s' completed.", tool_name)
            logger.info("    ↳ Content: %s\n", str(content)[:500])
            return result
        except Exception as e:
            logger.error("❌ [SERVER ERROR] Tool '%s' crashed: %s", tool_name, e)
            raise

mcp.add_middleware(ToolObserverMiddleware())

# ---------------------------------------------------------------------------
# CORS
# ---------------------------------------------------------------------------
middleware_config = [
    StarletteMiddleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_methods=["GET", "POST", "DELETE", "OPTIONS"],
        allow_headers=[
            "mcp-protocol-version",
            "mcp-session-id",
            "Authorization",
            "Content-Type",
        ],
        expose_headers=["mcp-session-id"],
    )
]

# ===================================================================
# Tool 1 — execute_python_code
# ===================================================================

@mcp.tool()
def execute_python_code(code: str) -> str:
    """Execute Python code in a sandboxed environment.

    Parameters
    ----------
    code : str
        The Python source to execute.

    Returns
    -------
    str
        stdout output, or error message if execution failed.
    """
    if not code or not code.strip():
        return "Error: code cannot be empty"

    if len(code) > 50_000:
        return "Error: code exceeds maximum length (50,000 chars)"

    # --- Memory limit via injected script ---
    # We inject the resource limit directly into the subprocess script.
    # Applying it in the main process would incorrectly limit the MCP server!
    memory_limit_code = f"""
try:
    import resource
    desired = {PYTHON_MAX_MEMORY_MB} * 1024 * 1024
    resource.setrlimit(resource.RLIMIT_AS, (desired, desired))
except Exception:
    pass
"""
    full_code = memory_limit_code.strip() + "\n\n" + code

    # --- Timeout via subprocess ---
    try:
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".py", delete=False
        ) as tmp:
            tmp.write(full_code)
            tmp_path = tmp.name

        proc = subprocess.run(
            [sys.executable, tmp_path],
            capture_output=True,
            text=True,
            timeout=PYTHON_EXEC_TIMEOUT,
        )
    except subprocess.TimeoutExpired:
        return f"Error: execution timed out after {PYTHON_EXEC_TIMEOUT}s"
    except Exception as e:
        return f"Error: {e}"
    finally:
        try:
            os.unlink(tmp_path)
        except OSError:
            pass

    # --- Output handling ---
    stdout = proc.stdout[:PYTHON_MAX_OUTPUT]
    stderr = proc.stderr[:PYTHON_MAX_OUTPUT]

    if proc.returncode != 0:
        return f"Exit code: {proc.returncode}\nStderr: {stderr}\nStdout: {stdout}"
    if stdout:
        return stdout
    return "Code executed successfully with no stdout output."

# ===================================================================
# Tool 2 — playwright_navigate
# ===================================================================

@mcp.tool()
async def playwright_navigate(url: str) -> str:
    """Navigate the persistent browser to a URL.

    Parameters
    ----------
    url : str
        Target URL.  Will be prefixed with https:// if missing scheme.

    Returns
    -------
    str
        Navigation result including page title, or error message.
    """
    try:
        # --- Validate URL ---
        parsed = urlparse(url)
        if not parsed.scheme or not parsed.netloc:
            if re.match(r"^[a-zA-Z0-9]", url):
                url = "https://" + url
            else:
                return f"Error: invalid URL '{url}'"

        page = await browser_manager.get_page()
        response = await page.goto(
            url, wait_until="domcontentloaded", timeout=PLAYWRIGHT_NAV_TIMEOUT
        )

        title = await page.title()
        status = response.status if response else "unknown"
        return f"Successfully loaded {url} (status {status})\nPage Title: {title}"
    except Exception as e:
        return f"Navigation error: {e}"

# ===================================================================
# Tool 3 — playwright_get_page_text
# ===================================================================

@mcp.tool()
async def playwright_get_page_text() -> str:
    """Retrieve the rendered page text from the current browser session.

    Returns
    -------
    str
        The innerText of document.body, or an error message.
    """
    try:
        page = browser_manager.page
        if not page:
            return "Error: no active browser session. Run playwright_navigate first."

        # CRITICAL FIX: page.evaluate does not support a `timeout` argument.
        # We must wrap the call in asyncio.wait_for instead.
        text_content = await asyncio.wait_for(
            page.evaluate("() => document.body.innerText"),
            timeout=PLAYWRIGHT_TEXT_TIMEOUT / 1000.0,
        )
        return f"=== CURRENT BROWSER TEXT SCREEN ===\n\n{text_content.strip()}"
    except asyncio.TimeoutError:
        return "Error: Page text extraction timed out."
    except Exception as e:
        return f"Extraction error: {e}"

# ===================================================================
# Tool 4 — extract_to_markdown
# ===================================================================

SUPPORTED_EXTENSIONS = {
    ".pdf", ".docx", ".pptx", ".xlsx", ".html", ".htm",
    ".md", ".txt", ".csv", ".json", ".xml", ".epub",
    ".odt", ".rtf", ".jpg", ".jpeg", ".png", ".gif",
    ".webp", ".bmp",
}

@mcp.tool()
def extract_to_markdown(relative_path: str) -> str:
    """Convert a file to Markdown using markitdown.

    Parameters
    ----------
    relative_path : str
        Path relative to WORKSPACE_ROOT.

    Returns
    -------
    str
        Markdown content, or an error message.
    """
    try:
        target = _resolve_safe_path(relative_path)
    except PermissionError as e:
        return f"Error: {e}"

    if not os.path.isfile(target):
        return f"Error: file not found at {relative_path}"

    # --- Size check ---
    file_size = os.path.getsize(target)
    if file_size > EXTRACT_MAX_FILE_SIZE:
        return f"Error: file size ({file_size} bytes) exceeds limit ({EXTRACT_MAX_FILE_SIZE} bytes)"

    # --- Extension check ---
    ext = os.path.splitext(target)[1].lower()
    if ext and ext not in SUPPORTED_EXTENSIONS:
        return f"Error: unsupported file type '{ext}'. Supported: {', '.join(sorted(SUPPORTED_EXTENSIONS))}"

    # --- Convert ---
    try:
        from markitdown import MarkItDown

        mid = MarkItDown()
        result = mid.convert(target)
        return result.text_content
    except ImportError:
        return "Error: markitdown library is not installed."
    except Exception as e:
        return f"Extraction error: {e}"

# ===================================================================
# Execution
# ===================================================================
if __name__ == "__main__":
    mcp.run(transport="http", host="0.0.0.0", port=8000, middleware=middleware_config)