#!/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 json
import logging
import asyncio
import subprocess
import tempfile
import secrets
import time
from typing import Dict, 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 FileResponse, JSONResponse

# ---------------------------------------------------------------------------
# Logging (JSON Format)
# ---------------------------------------------------------------------------
class JSONLogFormatter(logging.Formatter):
    """Custom JSON formatter for professional structured logging."""
    def format(self, record):
        log_record = {
            "timestamp": self.formatTime(record, self.datefmt),
            "level": record.levelname,
            "logger": record.name,
            "message": record.getMessage(),
            "file": f"{record.filename}:{record.lineno}"
        }
        if record.exc_info:
            log_record["exception"] = self.formatException(record.exc_info)
        return json.dumps(log_record)

json_handler = logging.StreamHandler(sys.stdout)
json_handler.setFormatter(JSONLogFormatter())

logging.basicConfig(
    level=logging.INFO,
    handlers=[json_handler],
    force=True  # Forces Uvicorn/FastAPI to also use this JSON handler
)
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

# Defaults to /workspace to match Docker volumes
WORKSPACE_MOUNT = os.path.abspath(os.environ.get("WORKSPACE_MOUNT", "/workspace"))

logger.info(f"Workspace Root configured to: {WORKSPACE_MOUNT}")

PUBLIC_BASE_URL = os.environ.get("PUBLIC_BASE_URL", "http://localhost:8000")
DOWNLOAD_DEFAULT_TTL = 10
DOWNLOAD_MIN_TTL = 1
DOWNLOAD_MAX_TTL = 60

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

TOKEN_STORE: Dict[str, dict] = {}

# ---------------------------------------------------------------------------
# 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…")
            
            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_MOUNT, blocking traversal."""
    target = os.path.abspath(os.path.join(WORKSPACE_MOUNT, relative_path))
    if not target.startswith(WORKSPACE_MOUNT):
        raise PermissionError("Access denied: path escapes WORKSPACE_MOUNT")
    return target


def _validate_download_path(raw_path: str) -> str:
    canonical_path = os.path.realpath(raw_path)
    if not canonical_path.startswith(WORKSPACE_MOUNT + "/"):
        raise PermissionError("Access denied: path escapes WORKSPACE_MOUNT")
    if not os.path.exists(canonical_path):
        raise FileNotFoundError(f"File not found: {raw_path}")
    if os.path.isdir(canonical_path):
        raise IsADirectoryError(
            "Target path is a directory. Please compress it to an archive first."
        )
    return canonical_path


def _purge_expired_tokens():
    now = time.time()
    expired_keys = [k for k, v in TOKEN_STORE.items() if now > v["expires_at"]]
    for k in expired_keys:
        TOKEN_STORE.pop(k, None)

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


@mcp.custom_route("/download", methods=["GET"])
async def download_file(request):
    _purge_expired_tokens()

    token = request.query_params.get("token")
    if not token:
        return JSONResponse({"detail": "Missing token parameter."}, status_code=400)

    record = TOKEN_STORE.get(token)
    if not record:
        return JSONResponse(
            {"detail": "Invalid or previously used download link."}, status_code=403
        )

    del TOKEN_STORE[token]

    if time.time() > record["expires_at"]:
        return JSONResponse({"detail": "Download link has expired."}, status_code=410)

    file_path = record["file_path"]
    if not os.path.exists(file_path):
        return JSONResponse(
            {"detail": "File no longer exists on host."}, status_code=404
        )

    return FileResponse(
        path=file_path,
        filename=os.path.basename(file_path),
        media_type="application/octet-stream",
    )

# ---------------------------------------------------------------------------
# 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(f"LLM REQUEST | Tool: '{tool_name}' | Args: {arguments}")
        try:
            result = await call_next(context)
            content = getattr(result, "content", result)
            preview = str(content)[:250].replace('\n', ' ') # Clean up newlines for JSON reading
            logger.info(f"SERVER RESPONSE | Tool: '{tool_name}' completed | Preview: {preview}")
            return result
        except Exception as e:
            logger.error(f"SERVER ERROR | Tool '{tool_name}' crashed: {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_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

    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

    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:
        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."

        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_MOUNT.

    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}"

    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)"

    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))}"

    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}"

# ===================================================================
# Tool 5 — generate_download_link
# ===================================================================

@mcp.tool()
def generate_download_link(path: str, ttl_minutes: int = DOWNLOAD_DEFAULT_TTL) -> str:
    """Generates a secure, single-use, temporary HTTPS download link for a file
    stored within the /workspace directory.

    Parameters
    ----------
    path : str
        The absolute path to the file inside the container (must reside within /workspace).
    ttl_minutes : int
        Validity period of the link in minutes. Default is 10. Clamped to 1–60.

    Returns
    -------
    str
        A clickable download URL, or an error message if validation fails.
    """
    try:
        abs_path = _validate_download_path(path)
    except (PermissionError, FileNotFoundError, IsADirectoryError) as e:
        return f"Error: {e}"

    _purge_expired_tokens()

    effective_ttl = max(DOWNLOAD_MIN_TTL, min(ttl_minutes, DOWNLOAD_MAX_TTL))
    token = secrets.token_urlsafe(32)
    expires_at = time.time() + (effective_ttl * 60)

    TOKEN_STORE[token] = {
        "file_path": abs_path,
        "expires_at": expires_at,
    }

    return f"{PUBLIC_BASE_URL}/download?token={token}"


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