#!/usr/bin/env python3
"""
MCP Server for serving local documents via SSE using FastMCP.
This server provides file reading capabilities for the localdocs directory.
"""

import os
import base64
import shutil
from pathlib import Path
from fastmcp import FastMCP
import sentry_sdk

sentry_sdk.init(
    dsn="https://fc33f2349050895a6cf1e2153a5d5f2c@sentry.eroomai.com/2",
    send_default_pii=True,
    traces_sample_rate=1.0,
)

from mcp_docparser.json_utils import safe_json_dumps

# Get docs directory from environment
DOCS_DIR = Path(os.environ.get("DOCS_DIR", "/home/eroom/Works/mcp-localdocs"))

# Create FastMCP server
mcp = FastMCP("mcp-localdocs")


# Binary file extensions that should be read as base64
BINARY_EXTENSIONS = {
    '.pdf', '.png', '.jpg', '.jpeg', '.gif', '.webp', '.bmp',
    '.zip', '.tar', '.gz', '.rar', '.7z',
    '.doc', '.docx', '.xls', '.xlsx', '.ppt', '.pptx',
    '.mp3', '.mp4', '.wav', '.avi', '.mov',
    '.bin', '.exe', '.dll', '.so', '.dylib',
}


def is_binary_file(file_path: Path) -> bool:
    """Check if a file should be treated as binary based on extension."""
    return file_path.suffix.lower() in BINARY_EXTENSIONS


def get_mime_type(file_path: Path) -> str:
    """Get MIME type based on file extension."""
    mime_types = {
        '.pdf': 'application/pdf',
        '.png': 'image/png',
        '.jpg': 'image/jpeg',
        '.jpeg': 'image/jpeg',
        '.gif': 'image/gif',
        '.webp': 'image/webp',
        '.bmp': 'image/bmp',
        '.zip': 'application/zip',
        '.tar': 'application/x-tar',
        '.gz': 'application/gzip',
        '.doc': 'application/msword',
        '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
        '.xls': 'application/vnd.ms-excel',
        '.xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
        '.ppt': 'application/vnd.ms-powerpoint',
        '.pptx': 'application/vnd.openxmlformats-officedocument.presentationml.presentation',
        '.mp3': 'audio/mpeg',
        '.mp4': 'video/mp4',
        '.wav': 'audio/wav',
        '.json': 'application/json',
        '.xml': 'application/xml',
        '.html': 'text/html',
        '.css': 'text/css',
        '.js': 'application/javascript',
        '.txt': 'text/plain',
        '.md': 'text/markdown',
    }
    return mime_types.get(file_path.suffix.lower(), 'application/octet-stream')


@mcp.tool()
def read_doc(doc_name: str) -> str:
    """Read content from a document in the localdocs directory.

    For text files, returns the content as plain text.
    For binary files (PDF, images, etc.), returns base64 encoded content with metadata.

    Args:
        doc_name: Name of the document to read (e.g., 'client_meeting.md', 'contract.pdf')
    """
    try:
        file_path = DOCS_DIR / doc_name

        # Security check
        if not file_path.resolve().is_relative_to(DOCS_DIR.resolve()):
            return f"Error: Access denied - path outside allowed directory"

        if not file_path.exists():
            return f"Error: Document not found: {doc_name}"

        # Check if binary file
        if is_binary_file(file_path):
            with open(file_path, 'rb') as f:
                content = f.read()

            mime_type = get_mime_type(file_path)
            encoded = base64.b64encode(content).decode('utf-8')

            # Return as JSON-like structure for easy parsing
            return safe_json_dumps({
                "type": "binary",
                "filename": doc_name,
                "mime_type": mime_type,
                "size": len(content),
                "content_base64": encoded
            })

        # Text file
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()

        return content
    except UnicodeDecodeError:
        # If text reading fails, try as binary
        try:
            with open(file_path, 'rb') as f:
                content = f.read()

            mime_type = get_mime_type(file_path)
            encoded = base64.b64encode(content).decode('utf-8')

            return safe_json_dumps({
                "type": "binary",
                "filename": doc_name,
                "mime_type": mime_type,
                "size": len(content),
                "content_base64": encoded
            })
        except Exception as e:
            return f"Error reading document: {str(e)}"
    except Exception as e:
        return f"Error reading document: {str(e)}"


@mcp.tool()
def read_binary_doc(doc_name: str) -> str:
    """Read a binary file (PDF, image, etc.) from the localdocs directory as base64.

    Always returns base64 encoded content regardless of file type.

    Args:
        doc_name: Name of the document to read (e.g., 'contract.pdf', 'receipt.png')

    Returns:
        JSON string with filename, mime_type, size, and content_base64
    """
    try:
        file_path = DOCS_DIR / doc_name

        # Security check
        if not file_path.resolve().is_relative_to(DOCS_DIR.resolve()):
            return f"Error: Access denied - path outside allowed directory"

        if not file_path.exists():
            return f"Error: Document not found: {doc_name}"

        with open(file_path, 'rb') as f:
            content = f.read()

        mime_type = get_mime_type(file_path)
        encoded = base64.b64encode(content).decode('utf-8')

        return safe_json_dumps({
            "type": "binary",
            "filename": doc_name,
            "mime_type": mime_type,
            "size": len(content),
            "content_base64": encoded
        })
    except Exception as e:
        return f"Error reading document: {str(e)}"


@mcp.tool()
def list_docs(pattern: str = "*", path: str = "", recursive: bool = False) -> str:
    """List all documents in the localdocs directory.

    Args:
        pattern: Optional glob pattern to filter documents (e.g., '*.json', '*.pdf')
        path: Relative path within localdocs to list from (default: root)
        recursive: If true, search in all subdirectories (default: false)
    """
    try:
        target_dir = DOCS_DIR / path if path else DOCS_DIR

        # Security check
        if not target_dir.resolve().is_relative_to(DOCS_DIR.resolve()):
            return f"Error: Access denied - path outside allowed directory"

        if not target_dir.exists():
            return f"Error: Directory not found: {path}"

        if not target_dir.is_dir():
            return f"Error: Path is not a directory: {path}"

        files = []
        glob_pattern = f"**/{pattern}" if recursive else pattern

        for file_path in target_dir.glob(glob_pattern):
            if file_path.is_file():
                rel_path = str(file_path.relative_to(DOCS_DIR))
                stat = file_path.stat()
                files.append({
                    "name": file_path.name,
                    "path": rel_path,
                    "size": stat.st_size,
                    "modified": stat.st_mtime,
                })

        # Sort by path
        files.sort(key=lambda x: x["path"])

        if not files:
            return safe_json_dumps({
                "path": path or "/",
                "pattern": pattern,
                "recursive": recursive,
                "files": [],
                "total": 0
            })

        return safe_json_dumps({
            "path": path or "/",
            "pattern": pattern,
            "recursive": recursive,
            "files": files,
            "total": len(files)
        })
    except Exception as e:
        return f"Error listing documents: {str(e)}"


@mcp.tool()
def write_file(path: str, content: str, overwrite: bool = True) -> str:
    """Write content to a file in the localdocs directory.

    Args:
        path: File path to write to
        content: Content to write
        overwrite: Whether to overwrite existing file (default: true)
    """
    try:
        file_path = DOCS_DIR / path

        # Security check
        if not file_path.resolve().is_relative_to(DOCS_DIR.resolve()):
            return f"Error: Access denied - path outside allowed directory"

        # Check if file exists
        if file_path.exists() and not overwrite:
            return f"Error: File already exists: {path}. Use overwrite=true to replace."

        # Create parent directories if needed
        file_path.parent.mkdir(parents=True, exist_ok=True)

        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(content)

        return f"Successfully wrote {len(content)} characters to {path}"
    except Exception as e:
        return f"Error writing file: {str(e)}"


@mcp.tool()
def list_folders(path: str = "", recursive: bool = False) -> str:
    """List all folders in the localdocs directory.

    Args:
        path: Relative path within localdocs to list folders from (default: root)
        recursive: If true, list all nested folders recursively (default: false)

    Returns:
        JSON string with folder structure including name, path, and file/folder counts
    """
    try:
        target_dir = DOCS_DIR / path if path else DOCS_DIR

        # Security check
        if not target_dir.resolve().is_relative_to(DOCS_DIR.resolve()):
            return f"Error: Access denied - path outside allowed directory"

        if not target_dir.exists():
            return f"Error: Directory not found: {path}"

        if not target_dir.is_dir():
            return f"Error: Path is not a directory: {path}"

        def get_folder_info(folder_path: Path) -> dict:
            """Get information about a folder."""
            rel_path = str(folder_path.relative_to(DOCS_DIR))
            if rel_path == ".":
                rel_path = ""

            files = [f for f in folder_path.iterdir() if f.is_file()]
            subdirs = [d for d in folder_path.iterdir() if d.is_dir()]

            return {
                "name": folder_path.name or "root",
                "path": rel_path,
                "file_count": len(files),
                "folder_count": len(subdirs),
            }

        def list_folders_recursive(folder_path: Path) -> list:
            """Recursively list all folders."""
            result = []
            for item in sorted(folder_path.iterdir()):
                if item.is_dir():
                    folder_info = get_folder_info(item)
                    folder_info["subfolders"] = list_folders_recursive(item)
                    result.append(folder_info)
            return result

        if recursive:
            root_info = get_folder_info(target_dir)
            root_info["subfolders"] = list_folders_recursive(target_dir)
            return safe_json_dumps(root_info)
        else:
            folders = []
            for item in sorted(target_dir.iterdir()):
                if item.is_dir():
                    folders.append(get_folder_info(item))

            return safe_json_dumps({
                "current_path": path or "/",
                "folders": folders,
                "total": len(folders)
            })
    except Exception as e:
        return f"Error listing folders: {str(e)}"


@mcp.tool()
def create_folder(path: str) -> str:
    """Create a new folder in the localdocs directory.

    Args:
        path: Relative path of the folder to create (e.g., 'projects/2024' or 'invoices')

    Returns:
        Success or error message
    """
    try:
        folder_path = DOCS_DIR / path

        # Security check
        if not folder_path.resolve().is_relative_to(DOCS_DIR.resolve()):
            return f"Error: Access denied - path outside allowed directory"

        if folder_path.exists():
            if folder_path.is_dir():
                return f"Folder already exists: {path}"
            else:
                return f"Error: A file with that name already exists: {path}"

        folder_path.mkdir(parents=True, exist_ok=True)
        return f"Successfully created folder: {path}"
    except Exception as e:
        return f"Error creating folder: {str(e)}"


@mcp.tool()
def delete_folder(path: str, force: bool = False) -> str:
    """Delete a folder from the localdocs directory.

    Args:
        path: Relative path of the folder to delete
        force: If true, delete non-empty folders recursively (default: false)

    Returns:
        Success or error message
    """
    try:
        folder_path = DOCS_DIR / path

        # Security check
        if not folder_path.resolve().is_relative_to(DOCS_DIR.resolve()):
            return f"Error: Access denied - path outside allowed directory"

        # Prevent deleting root
        if folder_path.resolve() == DOCS_DIR.resolve():
            return "Error: Cannot delete the root directory"

        if not folder_path.exists():
            return f"Error: Folder not found: {path}"

        if not folder_path.is_dir():
            return f"Error: Path is not a folder: {path}"

        # Check if folder is empty
        contents = list(folder_path.iterdir())
        if contents and not force:
            return f"Error: Folder is not empty ({len(contents)} items). Use force=true to delete recursively."

        if force:
            shutil.rmtree(folder_path)
        else:
            folder_path.rmdir()

        return f"Successfully deleted folder: {path}"
    except Exception as e:
        return f"Error deleting folder: {str(e)}"


@mcp.tool()
def move_item(source: str, destination: str) -> str:
    """Move or rename a file or folder within the localdocs directory.

    Args:
        source: Relative path of the file/folder to move
        destination: New relative path (can be a new name or different location)

    Returns:
        Success or error message
    """
    try:
        source_path = DOCS_DIR / source
        dest_path = DOCS_DIR / destination

        # Security checks
        if not source_path.resolve().is_relative_to(DOCS_DIR.resolve()):
            return f"Error: Access denied - source path outside allowed directory"

        if not dest_path.resolve().is_relative_to(DOCS_DIR.resolve()):
            return f"Error: Access denied - destination path outside allowed directory"

        if not source_path.exists():
            return f"Error: Source not found: {source}"

        if dest_path.exists():
            return f"Error: Destination already exists: {destination}"

        # Create parent directories if needed
        dest_path.parent.mkdir(parents=True, exist_ok=True)

        shutil.move(str(source_path), str(dest_path))

        item_type = "folder" if dest_path.is_dir() else "file"
        return f"Successfully moved {item_type}: {source} -> {destination}"
    except Exception as e:
        return f"Error moving item: {str(e)}"


@mcp.tool()
def copy_item(source: str, destination: str) -> str:
    """Copy a file or folder within the localdocs directory.

    Args:
        source: Relative path of the file/folder to copy
        destination: Destination relative path

    Returns:
        Success or error message
    """
    try:
        source_path = DOCS_DIR / source
        dest_path = DOCS_DIR / destination

        # Security checks
        if not source_path.resolve().is_relative_to(DOCS_DIR.resolve()):
            return f"Error: Access denied - source path outside allowed directory"

        if not dest_path.resolve().is_relative_to(DOCS_DIR.resolve()):
            return f"Error: Access denied - destination path outside allowed directory"

        if not source_path.exists():
            return f"Error: Source not found: {source}"

        if dest_path.exists():
            return f"Error: Destination already exists: {destination}"

        # Create parent directories if needed
        dest_path.parent.mkdir(parents=True, exist_ok=True)

        if source_path.is_dir():
            shutil.copytree(str(source_path), str(dest_path))
            return f"Successfully copied folder: {source} -> {destination}"
        else:
            shutil.copy2(str(source_path), str(dest_path))
            return f"Successfully copied file: {source} -> {destination}"
    except Exception as e:
        return f"Error copying item: {str(e)}"


@mcp.tool()
def delete_file(path: str) -> str:
    """Delete a file from the localdocs directory.

    Args:
        path: Relative path of the file to delete

    Returns:
        Success or error message
    """
    try:
        file_path = DOCS_DIR / path

        # Security check
        if not file_path.resolve().is_relative_to(DOCS_DIR.resolve()):
            return f"Error: Access denied - path outside allowed directory"

        if not file_path.exists():
            return f"Error: File not found: {path}"

        if file_path.is_dir():
            return f"Error: Path is a directory, not a file. Use delete_folder instead: {path}"

        file_path.unlink()
        return f"Successfully deleted file: {path}"
    except Exception as e:
        return f"Error deleting file: {str(e)}"


if __name__ == "__main__":
    print(f"Starting MCP LocalDocs Server with FastMCP...")
    print(f"Serving files from: {DOCS_DIR}")
    print(f"HTTP endpoint: http://localhost:8012/mcp")
    # Use streamable-http instead of SSE to avoid MemoryObjectReceiveStream leak
    mcp.run(transport="streamable-http", host="0.0.0.0", port=8012)