import subprocess
import shlex
import os
from pathlib import Path
from typing import List, Optional

# In a real MCP environment, the 'mcp' object with its decorators
# would be provided by the framework. This is a placeholder for standalone validation.
class _MCP:
    def tool(self, *args, **kwargs):
        def decorator(f):
            return f
        return decorator
mcp = _MCP()

### Bamtools Tool Definitions ###

from mcp.server.fastmcp import FastMCP

SERVER_NAME = 'local_bamtools'
mcp = FastMCP(SERVER_NAME)

@mcp.tool()
def convert(
    in_bam: Path,
    format: str,
    out: Optional[Path] = None,
    region: Optional[str] = None,
) -> dict:
    """
    Converts a BAM file to various other formats.

    Args:
        in_bam: The input BAM file.
        format: The output format. Must be one of: bed, fasta, fastq, json, pileup, sam, yaml.
        out: The output filename. If not provided, output is sent to stdout.
        region: Only convert alignments that overlap this region (e.g., "ref:start-end").
    
    Returns:
        A dictionary containing the command executed, stdout, stderr, and output files.
    """
    if not in_bam.is_file():
        raise FileNotFoundError(f"Input BAM file not found: {in_bam}")

    valid_formats = {"bed", "fasta", "fastq", "json", "pileup", "sam", "yaml"}
    if format not in valid_formats:
        raise ValueError(f"Invalid format '{format}'. Must be one of {valid_formats}")

    cmd = ["bamtools", "convert", "-in", str(in_bam), "-format", format]
    output_files = []

    if out:
        cmd.extend(["-out", str(out)])
        output_files.append(str(out))
    if region:
        cmd.extend(["-region", region])

    command_executed = shlex.join(cmd)
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        return {
            "command_executed": command_executed,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": output_files,
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout,
            "stderr": e.stderr,
            "error": f"bamtools convert failed with exit code {e.returncode}",
            "output_files": [],
        }

@mcp.tool()
def count(in_bams: List[Path]) -> dict:
    """
    Prints the number of alignments in one or more BAM files.

    Args:
        in_bams: A list of input BAM files.
    
    Returns:
        A dictionary containing the command executed, stdout (the count), and stderr.
    """
    if not in_bams:
        raise ValueError("At least one input BAM file must be provided.")

    cmd = ["bamtools", "count"]
    for bam_file in in_bams:
        if not bam_file.is_file():
            raise FileNotFoundError(f"Input BAM file not found: {bam_file}")
        cmd.extend(["-in", str(bam_file)])

    command_executed = shlex.join(cmd)
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        return {
            "command_executed": command_executed,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": [],
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout,
            "stderr": e.stderr,
            "error": f"bamtools count failed with exit code {e.returncode}",
            "output_files": [],
        }

@mcp.tool()
def coverage(in_bams: List[Path]) -> dict:
    """
    Prints coverage statistics from one or more BAM files.

    Args:
        in_bams: A list of input BAM files.
    
    Returns:
        A dictionary containing the command executed, stdout (coverage stats), and stderr.
    """
    if not in_bams:
        raise ValueError("At least one input BAM file must be provided.")

    cmd = ["bamtools", "coverage"]
    for bam_file in in_bams:
        if not bam_file.is_file():
            raise FileNotFoundError(f"Input BAM file not found: {bam_file}")
        cmd.extend(["-in", str(bam_file)])

    command_executed = shlex.join(cmd)
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        return {
            "command_executed": command_executed,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": [],
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout,
            "stderr": e.stderr,
            "error": f"bamtools coverage failed with exit code {e.returncode}",
            "output_files": [],
        }

@mcp.tool()
def filter_bam(
    in_bams: List[Path],
    out_bam: Optional[Path] = None,
    region: Optional[str] = None,
    property_filters: Optional[List[str]] = None,
    tag_filters: Optional[List[str]] = None,
    script: Optional[Path] = None,
    list_file: Optional[Path] = None,
    force_compression: bool = False,
) -> dict:
    """
    Filters BAM file(s) by user-specified criteria.

    Args:
        in_bams: The input BAM file(s).
        out_bam: The output BAM file. If not provided, output is sent to stdout.
        region: Only keeps alignments that overlap this region.
        property_filters: List of "TAG:VALUE" strings to filter by BamAlignment property.
        tag_filters: List of "TAG:VALUE" strings to filter by BamAlignment tag.
        script: A file containing a list of filter rules.
        list_file: A file containing a list of read names to be kept.
        force_compression: If True, the output BAM file will be compressed.
    
    Returns:
        A dictionary containing the command executed, stdout, stderr, and output files.
    """
    if not in_bams:
        raise ValueError("At least one input BAM file must be provided.")

    cmd = ["bamtools", "filter"]
    output_files = []

    for bam_file in in_bams:
        if not bam_file.is_file():
            raise FileNotFoundError(f"Input BAM file not found: {bam_file}")
        cmd.extend(["-in", str(bam_file)])

    if out_bam:
        cmd.extend(["-out", str(out_bam)])
        output_files.append(str(out_bam))
    if region:
        cmd.extend(["-region", region])
    if property_filters:
        for p_filter in property_filters:
            cmd.extend(["-property", p_filter])
    if tag_filters:
        for t_filter in tag_filters:
            cmd.extend(["-tag", t_filter])
    if script:
        if not script.is_file():
            raise FileNotFoundError(f"Script file not found: {script}")
        cmd.extend(["-script", str(script)])
    if list_file:
        if not list_file.is_file():
            raise FileNotFoundError(f"List file not found: {list_file}")
        cmd.extend(["-list", str(list_file)])
    if force_compression:
        cmd.append("-forceCompression")

    command_executed = shlex.join(cmd)
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        return {
            "command_executed": command_executed,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": output_files,
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout,
            "stderr": e.stderr,
            "error": f"bamtools filter failed with exit code {e.returncode}",
            "output_files": [],
        }

@mcp.tool()
def header(in_bam: Path) -> dict:
    """
    Prints the header from a BAM file.

    Args:
        in_bam: The input BAM file.
    
    Returns:
        A dictionary containing the command executed, stdout (the header), and stderr.
    """
    if not in_bam.is_file():
        raise FileNotFoundError(f"Input BAM file not found: {in_bam}")

    cmd = ["bamtools", "header", "-in", str(in_bam)]
    command_executed = shlex.join(cmd)
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        return {
            "command_executed": command_executed,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": [],
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout,
            "stderr": e.stderr,
            "error": f"bamtools header failed with exit code {e.returncode}",
            "output_files": [],
        }

@mcp.tool()
def index(in_bam: Path) -> dict:
    """
    Generates an index for a BAM file. The index file (.bai) is created in the same directory.

    Args:
        in_bam: The input BAM file to index.
    
    Returns:
        A dictionary containing the command executed, stdout, stderr, and the path to the generated index file.
    """
    if not in_bam.is_file():
        raise FileNotFoundError(f"Input BAM file not found: {in_bam}")

    cmd = ["bamtools", "index", "-in", str(in_bam)]
    output_index_file = in_bam.with_suffix(in_bam.suffix + '.bai')
    
    command_executed = shlex.join(cmd)
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        if not output_index_file.is_file():
             raise FileNotFoundError(f"Expected index file was not created: {output_index_file}")
        return {
            "command_executed": command_executed,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": [str(output_index_file)],
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout,
            "stderr": e.stderr,
            "error": f"bamtools index failed with exit code {e.returncode}",
            "output_files": [],
        }

@mcp.tool()
def merge(
    in_bams: List[Path],
    out_bam: Optional[Path] = None,
    region: Optional[str] = None,
    force: bool = False,
) -> dict:
    """
    Merges multiple BAM files into a single file.

    Args:
        in_bams: A list of input BAM files to merge. Must contain at least two files.
        out_bam: The output BAM file. If not provided, output is sent to stdout.
        region: Merges only alignments that overlap this region.
        force: Forces merge, even if headers are not identical. Uses the first BAM file's header.
    
    Returns:
        A dictionary containing the command executed, stdout, stderr, and output files.
    """
    if len(in_bams) < 2:
        raise ValueError("At least two input BAM files must be provided for merging.")

    cmd = ["bamtools", "merge"]
    output_files = []

    for bam_file in in_bams:
        if not bam_file.is_file():
            raise FileNotFoundError(f"Input BAM file not found: {bam_file}")
        cmd.extend(["-in", str(bam_file)])

    if out_bam:
        cmd.extend(["-out", str(out_bam)])
        output_files.append(str(out_bam))
    if region:
        cmd.extend(["-region", region])
    if force:
        cmd.append("-force")

    command_executed = shlex.join(cmd)
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        return {
            "command_executed": command_executed,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": output_files,
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout,
            "stderr": e.stderr,
            "error": f"bamtools merge failed with exit code {e.returncode}",
            "output_files": [],
        }

@mcp.tool()
def random(
    in_bams: List[Path],
    out_bam: Optional[Path] = None,
    n: int = 100,
    seed: int = 0,
) -> dict:
    """
    Selects a random number of alignments from existing BAM file(s).

    Args:
        in_bams: The input BAM file(s).
        out_bam: The output BAM file. If not provided, output is sent to stdout.
        n: Number of alignments to select.
        seed: Random seed.
    
    Returns:
        A dictionary containing the command executed, stdout, stderr, and output files.
    """
    if not in_bams:
        raise ValueError("At least one input BAM file must be provided.")

    cmd = ["bamtools", "random", "-n", str(n), "-seed", str(seed)]
    output_files = []

    for bam_file in in_bams:
        if not bam_file.is_file():
            raise FileNotFoundError(f"Input BAM file not found: {bam_file}")
        cmd.extend(["-in", str(bam_file)])

    if out_bam:
        cmd.extend(["-out", str(out_bam)])
        output_files.append(str(out_bam))

    command_executed = shlex.join(cmd)
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        return {
            "command_executed": command_executed,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": output_files,
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout,
            "stderr": e.stderr,
            "error": f"bamtools random failed with exit code {e.returncode}",
            "output_files": [],
        }

@mcp.tool()
def resolve(in_bam: Path, out_bam: Path) -> dict:
    """
    Resolves paired-end reads, marking the IsProperPair flag as needed.

    Args:
        in_bam: The input BAM file.
        out_bam: The output BAM file.
    
    Returns:
        A dictionary containing the command executed, stdout, stderr, and the output file path.
    """
    if not in_bam.is_file():
        raise FileNotFoundError(f"Input BAM file not found: {in_bam}")

    cmd = ["bamtools", "resolve", "-in", str(in_bam), "-out", str(out_bam)]
    command_executed = shlex.join(cmd)
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        return {
            "command_executed": command_executed,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": [str(out_bam)],
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout,
            "stderr": e.stderr,
            "error": f"bamtools resolve failed with exit code {e.returncode}",
            "output_files": [],
        }

@mcp.tool()
def revert(in_bam: Path, out_bam: Path) -> dict:
    """
    Removes duplicate marks and restores original base qualities from a BAM file.

    Args:
        in_bam: The input BAM file.
        out_bam: The output BAM file.
    
    Returns:
        A dictionary containing the command executed, stdout, stderr, and the output file path.
    """
    if not in_bam.is_file():
        raise FileNotFoundError(f"Input BAM file not found: {in_bam}")

    cmd = ["bamtools", "revert", "-in", str(in_bam), "-out", str(out_bam)]
    command_executed = shlex.join(cmd)
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        return {
            "command_executed": command_executed,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": [str(out_bam)],
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout,
            "stderr": e.stderr,
            "error": f"bamtools revert failed with exit code {e.returncode}",
            "output_files": [],
        }

@mcp.tool()
def sort(
    in_bam: Path,
    out_bam: Optional[Path] = None,
    by: str = "position",
    order: str = "ascending",
    n: int = 1000,
    mem: int = 512,
) -> dict:
    """
    Sorts a BAM file by position or name.

    Args:
        in_bam: The input BAM file.
        out_bam: The output BAM file. If not provided, output is sent to stdout.
        by: Sort by 'position' or 'name'.
        order: Sort in 'ascending' or 'descending' order.
        n: Max number of alignments to buffer in memory.
        mem: Max memory to use (in megabytes).
    
    Returns:
        A dictionary containing the command executed, stdout, stderr, and output files.
    """
    if not in_bam.is_file():
        raise FileNotFoundError(f"Input BAM file not found: {in_bam}")
    if by not in ["position", "name"]:
        raise ValueError("Sort 'by' must be 'position' or 'name'.")
    if order not in ["ascending", "descending"]:
        raise ValueError("Sort 'order' must be 'ascending' or 'descending'.")

    cmd = ["bamtools", "sort", "-in", str(in_bam), "-by", by, "-order", order, "-n", str(n), "-mem", str(mem)]
    output_files = []

    if out_bam:
        cmd.extend(["-out", str(out_bam)])
        output_files.append(str(out_bam))

    command_executed = shlex.join(cmd)
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        return {
            "command_executed": command_executed,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": output_files,
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout,
            "stderr": e.stderr,
            "error": f"bamtools sort failed with exit code {e.returncode}",
            "output_files": [],
        }

@mcp.tool()
def split(
    in_bam: Path,
    output_dir: Path,
    stub: Optional[str] = None,
    by: Optional[str] = None,
    reference: bool = False,
    mapped: bool = False,
    paired: bool = False,
    tag: Optional[str] = None,
) -> dict:
    """
    Splits a BAM file based on a specified property, creating multiple output BAM files.

    Args:
        in_bam: The input BAM file.
        output_dir: The directory where output files will be created.
        stub: The prefix for output file names. Defaults to the input file name.
        by: Split file by a specific property.
        reference: Split by reference sequence.
        mapped: Split into mapped and unmapped files.
        paired: Split into paired and single-end files.
        tag: Split by the value of a specific tag.
    
    Returns:
        A dictionary containing the command executed, stdout, stderr, and a list of created files.
    """
    if not in_bam.is_file():
        raise FileNotFoundError(f"Input BAM file not found: {in_bam}")

    split_options = [by is not None, reference, mapped, paired, tag is not None]
    if sum(split_options) != 1:
        raise ValueError("Exactly one split criterion (by, reference, mapped, paired, or tag) must be specified.")

    output_dir.mkdir(parents=True, exist_ok=True)

    cmd = ["bamtools", "split", "-in", str(in_bam)]
    
    # Use the stub to direct output to the specified directory
    output_stub = stub if stub else in_bam.stem
    cmd.extend(["-stub", str(output_dir / output_stub)])

    if by:
        cmd.extend(["-by", by])
    if reference:
        cmd.append("-reference")
    if mapped:
        cmd.append("-mapped")
    if paired:
        cmd.append("-paired")
    if tag:
        cmd.extend(["-tag", tag])

    command_executed = shlex.join(cmd)
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        # Find all files created in the output directory
        created_files = [str(p) for p in output_dir.glob(f"{output_stub}*")]
        return {
            "command_executed": command_executed,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": created_files,
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout,
            "stderr": e.stderr,
            "error": f"bamtools split failed with exit code {e.returncode}",
            "output_files": [],
        }

@mcp.tool()
def stats(in_bams: List[Path], insert: bool = False) -> dict:
    """
    Prints basic statistics from one or more BAM files.

    Args:
        in_bams: A list of input BAM files.
        insert: If True, print insert size summary.
    
    Returns:
        A dictionary containing the command executed, stdout (the stats), and stderr.
    """
    if not in_bams:
        raise ValueError("At least one input BAM file must be provided.")

    cmd = ["bamtools", "stats"]
    for bam_file in in_bams:
        if not bam_file.is_file():
            raise FileNotFoundError(f"Input BAM file not found: {bam_file}")
        cmd.extend(["-in", str(bam_file)])

    if insert:
        cmd.append("-insert")

    command_executed = shlex.join(cmd)
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        return {
            "command_executed": command_executed,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": [],
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout,
            "stderr": e.stderr,
            "error": f"bamtools stats failed with exit code {e.returncode}",
            "output_files": [],
        }

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