import subprocess
import tempfile
from pathlib import Path
from typing import List, Optional

# MCP-standardized error handling
class MCPError(Exception):
    def __init__(self, message, stderr):
        self.message = message
        self.stderr = stderr
        super().__init__(self.message)

def _run_command(cmd: List[str]):
    """Helper function to run a command and handle exceptions."""
    try:
        command_str = " ".join(cmd)
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        return {
            "command_executed": command_str,
            "stdout": result.stdout,
            "stderr": result.stderr,
        }
    except FileNotFoundError:
        raise MCPError(f"Command '{cmd[0]}' not found. Is arvados-python-client installed and in the PATH?", "")
    except subprocess.CalledProcessError as e:
        raise MCPError(f"Command '{' '.join(cmd)}' failed with exit code {e.returncode}", e.stderr)

from mcp.server.fastmcp import FastMCP

SERVER_NAME = 'local_arvados_python_client'
mcp = FastMCP(SERVER_NAME)

@mcp.tool()
def arv_get(
    object_locator: str,
    destination_path: Optional[Path] = None,
    filename: Optional[str] = None,
    force: bool = False,
    recursive: bool = False,
    progress: bool = False,
) -> dict:
    """
    Downloads an object from Arvados Keep to a local path or standard output.

    Args:
        object_locator: The UUID or portable data hash (PDH) of the object to download.
        destination_path: Local path to save the object. If a directory, the original filename is used. If not provided, content is written to stdout.
        filename: The name of the file to get from a collection.
        force: Overwrite existing files if they exist.
        recursive: Recursively download collections.
        progress: Show a progress bar during download.

    Returns:
        A dictionary containing the command executed, stdout, stderr, and a list of output files.
    """
    cmd = ["arv-get"]

    if force:
        cmd.append("--force")
    if recursive:
        cmd.append("--recursive")
    if progress:
        cmd.append("--progress")
    if filename:
        cmd.extend(["--filename", filename])

    cmd.append(object_locator)

    output_files = []
    if destination_path:
        # Ensure parent directory exists if a full path is given
        if destination_path.parent and not destination_path.parent.is_dir():
            destination_path.parent.mkdir(parents=True, exist_ok=True)
        cmd.append(str(destination_path))
        output_files.append(str(destination_path))

    try:
        result = _run_command(cmd)
        result["output_files"] = output_files
        return result
    except MCPError as e:
        return {
            "command_executed": " ".join(cmd),
            "stdout": "",
            "stderr": e.stderr,
            "error": e.message,
            "output_files": [],
        }

@mcp.tool()
def arv_put(
    paths: List[Path],
    name: Optional[str] = None,
    owner: Optional[str] = None,
    portable_data_hash: bool = False,
    storage_class: Optional[str] = None,
    num_retries: int = 5,
    progress: bool = False,
) -> dict:
    """
    Uploads local files or directories to Arvados Keep.

    Args:
        paths: One or more local file or directory paths to upload.
        name: Set the name of the new collection.
        owner: Set the owner project UUID for the new collection.
        portable_data_hash: Use portable data hash to identify the collection.
        storage_class: The storage class for the data (e.g., 'default', 'trash').
        num_retries: Number of times to retry a failed API call.
        progress: Show a progress bar during upload.

    Returns:
        A dictionary containing the command executed, stdout, and stderr.
    """
    if not paths:
        raise ValueError("At least one path must be provided for upload.")

    for p in paths:
        if not p.exists():
            raise FileNotFoundError(f"Input path does not exist: {p}")

    cmd = ["arv-put"]

    if name:
        cmd.extend(["--name", name])
    if owner:
        cmd.extend(["--owner", owner])
    if portable_data_hash:
        cmd.append("--portable-data-hash")
    if storage_class:
        cmd.extend(["--storage-class", storage_class])
    if progress:
        cmd.append("--progress")
    
    # num_retries is a standard arvados client option
    cmd.extend(["--num-retries", str(num_retries)])

    cmd.extend([str(p) for p in paths])

    try:
        result = _run_command(cmd)
        result["output_files"] = []
        return result
    except MCPError as e:
        return {
            "command_executed": " ".join(cmd),
            "stdout": "",
            "stderr": e.stderr,
            "error": e.message,
            "output_files": [],
        }

@mcp.tool()
def arv_ls(
    object_locator: str,
    path: Optional[str] = None,
    long_format: bool = False,
    recursive: bool = False,
    human_readable: bool = False,
) -> dict:
    """
    Lists the contents of a collection in Arvados Keep.

    Args:
        object_locator: The UUID or portable data hash (PDH) of the collection to list.
        path: A path within the collection to list.
        long_format: Use a long listing format (similar to ls -l).
        recursive: List subdirectories recursively.
        human_readable: With long_format, print sizes in human-readable format (e.g., 1K, 234M, 2G).

    Returns:
        A dictionary containing the command executed, stdout (the listing), and stderr.
    """
    cmd = ["arv-ls"]

    if long_format:
        cmd.append("-l")
    if recursive:
        cmd.append("-R")
    if human_readable:
        cmd.append("-h")

    cmd.append(object_locator)
    if path:
        cmd.append(path)

    try:
        result = _run_command(cmd)
        result["output_files"] = []
        return result
    except MCPError as e:
        return {
            "command_executed": " ".join(cmd),
            "stdout": "",
            "stderr": e.stderr,
            "error": e.message,
            "output_files": [],
        }

@mcp.tool()
def arv_copy(
    src_locators: List[str],
    dest_project_uuid: str,
    name: Optional[str] = None,
    num_retries: int = 5,
) -> dict:
    """
    Copies one or more objects to a destination project in Arvados.

    Args:
        src_locators: A list of source object UUIDs or portable data hashes (PDHs).
        dest_project_uuid: The UUID of the destination project.
        name: New name for the object(s). If copying multiple objects, this is ignored.
        num_retries: Number of times to retry a failed API call.

    Returns:
        A dictionary containing the command executed, stdout, and stderr.
    """
    if not src_locators:
        raise ValueError("At least one source locator must be provided.")
    if not dest_project_uuid:
        raise ValueError("A destination project UUID must be provided.")

    cmd = ["arv-copy"]

    if name and len(src_locators) == 1:
        cmd.extend(["--name", name])
    
    cmd.extend(["--num-retries", str(num_retries)])
    
    cmd.extend(src_locators)
    cmd.append(dest_project_uuid)

    try:
        result = _run_command(cmd)
        result["output_files"] = []
        return result
    except MCPError as e:
        return {
            "command_executed": " ".join(cmd),
            "stdout": "",
            "stderr": e.stderr,
            "error": e.message,
            "output_files": [],
        }

@mcp.tool()
def arv_tag(
    locator: str,
    tags_to_add: Optional[List[str]] = None,
    keys_to_remove: Optional[List[str]] = None,
    list_tags: bool = False,
    num_retries: int = 5,
) -> dict:
    """
    Manipulates tags on an Arvados object. It can list, add, or remove tags.

    Args:
        locator: The UUID or portable data hash (PDH) of the object to tag.
        tags_to_add: A list of tags to add, in 'key' or 'key=value' format.
        keys_to_remove: A list of tag keys to remove.
        list_tags: If True, lists the existing tags on the object.
        num_retries: Number of times to retry a failed API call.

    Returns:
        A dictionary containing the command executed, stdout, and stderr.
    """
    cmd = ["arv-tag"]
    
    cmd.extend(["--num-retries", str(num_retries)])

    if list_tags:
        cmd.append("--list")
    
    cmd.append(locator)

    if tags_to_add:
        for tag in tags_to_add:
            if "=" not in tag and " " in tag:
                raise ValueError(f"Invalid tag format: '{tag}'. Use 'key=value' or 'key'.")
        cmd.extend(tags_to_add)

    if keys_to_remove:
        for key in keys_to_remove:
            cmd.extend(["--remove", key])

    try:
        result = _run_command(cmd)
        result["output_files"] = []
        return result
    except MCPError as e:
        return {
            "command_executed": " ".join(cmd),
            "stdout": "",
            "stderr": e.stderr,
            "error": e.message,
            "output_files": [],
        }

if __name__ == "__main__":
    mcp.run(transport="stdio")