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

from mcp.server.fastmcp import FastMCP

SERVER_NAME = 'local_bactopia'
mcp = FastMCP(SERVER_NAME)

@mcp.tool()
def bactopia_run(
    sample: str,
    outdir: str,
    datasets: str,
    r1: Optional[str] = None,
    r2: Optional[str] = None,
    se: Optional[str] = None,
    fastqs: Optional[str] = None,
    accession: Optional[str] = None,
    accessions: Optional[str] = None,
    profile: str = "conda",
    cpus: int = 2,
    max_memory: int = 8,
    resume: bool = False,
) -> dict:
    """
    Run the main Bactopia analysis pipeline for a bacterial genome.
    
    Args:
        sample: Name of the sample.
        outdir: Directory to store output results.
        datasets: Path to the Bactopia datasets directory.
        r1: Path to forward reads (Paired-end).
        r2: Path to reverse reads (Paired-end).
        se: Path to single-end reads.
        fastqs: Path to a FOFN (file-of-filenames) for multiple samples.
        accession: A single ENA/SRA accession (e.g., SRX000000).
        accessions: A file containing a list of ENA/SRA accessions.
        profile: Nextflow profile to use (e.g., conda, docker, singularity).
        cpus: Number of CPUs to allocate.
        max_memory: Maximum memory in GB to allocate.
        resume: Whether to resume a previous run.
    """
    # Input validation
    out_path = Path(outdir)
    ds_path = Path(datasets)
    
    if not ds_path.exists():
        return {"error": f"Datasets directory not found at {datasets}"}

    cmd = ["bactopia", "--sample", sample, "--outdir", str(out_path), "--datasets", str(ds_path)]
    
    # Input source logic
    if r1 and r2:
        if not Path(r1).exists() or not Path(r2).exists():
            return {"error": "R1 or R2 file does not exist"}
        cmd += ["--R1", r1, "--R2", r2]
    elif se:
        if not Path(se).exists():
            return {"error": "Single-end file does not exist"}
        cmd += ["--SE", se]
    elif fastqs:
        if not Path(fastqs).exists():
            return {"error": "Fastqs list file does not exist"}
        cmd += ["--fastqs", fastqs]
    elif accession:
        cmd += ["--accession", accession]
    elif accessions:
        if not Path(accessions).exists():
            return {"error": "Accessions list file does not exist"}
        cmd += ["--accessions", accessions]
    else:
        return {"error": "No input source provided (R1/R2, SE, fastqs, accession, or accessions required)"}

    # Performance and Nextflow options
    cmd += ["-profile", profile]
    cmd += ["--max_cpus", str(cpus)]
    cmd += ["--max_memory", f"{max_memory}.GB"]
    
    if resume:
        cmd.append("-resume")

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        return {
            "command_executed": " ".join(cmd),
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_directory": str(out_path.absolute())
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": " ".join(cmd),
            "error": str(e),
            "stdout": e.stdout,
            "stderr": e.stderr
        }

@mcp.tool()
def bactopia_prepare(
    path: str,
    recursive: bool = False,
    extension: str = ".fastq.gz",
    output_file: Optional[str] = None
) -> dict:
    """
    Prepare a FASTQ list (FOFN) for Bactopia from a directory of sequencing files.
    
    Args:
        path: Directory containing FASTQ files.
        recursive: Search for FASTQs in subdirectories.
        extension: File extension to look for.
        output_file: Optional path to save the generated list.
    """
    input_path = Path(path)
    if not input_path.is_dir():
        return {"error": f"Path {path} is not a directory"}

    cmd = ["bactopia", "prepare", str(input_path), "--extension", extension]
    if recursive:
        cmd.append("--recursive")

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        
        if output_file:
            out_p = Path(output_file)
            out_p.write_text(result.stdout)
            return {
                "command_executed": " ".join(cmd),
                "stdout": "File list generated and saved.",
                "output_file": str(out_p.absolute())
            }
            
        return {
            "command_executed": " ".join(cmd),
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": " ".join(cmd),
            "error": str(e),
            "stdout": e.stdout,
            "stderr": e.stderr
        }

@mcp.tool()
def bactopia_search(
    query: str,
    limit: int = 10,
    output_file: Optional[str] = None
) -> dict:
    """
    Search for bacterial genome accessions in ENA/SRA.
    
    Args:
        query: Search query (e.g., "Staphylococcus aureus").
        limit: Maximum number of results to return.
        output_file: Optional path to save the accession list.
    """
    cmd = ["bactopia", "search", query, "--limit", str(limit)]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        
        if output_file:
            out_p = Path(output_file)
            out_p.write_text(result.stdout)
            return {
                "command_executed": " ".join(cmd),
                "stdout": f"Search results saved to {output_file}",
                "output_file": str(out_p.absolute())
            }

        return {
            "command_executed": " ".join(cmd),
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": " ".join(cmd),
            "error": str(e),
            "stdout": e.stdout,
            "stderr": e.stderr
        }

@mcp.tool()
def bactopia_datasets(
    outdir: str,
    species: Optional[str] = None,
    include_genus: bool = False,
    limit: int = 100
) -> dict:
    """
    Download and setup Bactopia datasets/databases.
    
    Args:
        outdir: Directory to save the datasets.
        species: Optional species name to download specific datasets (e.g., "Salmonella enterica").
        include_genus: Include genus-level datasets.
        limit: Limit the number of genomes used for building datasets.
    """
    out_path = Path(outdir)
    cmd = ["bactopia", "datasets", "--outdir", str(out_path)]
    
    if species:
        cmd += ["--species", species]
    if include_genus:
        cmd.append("--include_genus")
    cmd += ["--limit", str(limit)]

    try:
        # This can be a long-running process
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        return {
            "command_executed": " ".join(cmd),
            "stdout": result.stdout,
            "stderr": result.stderr,
            "datasets_path": str(out_path.absolute())
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": " ".join(cmd),
            "error": str(e),
            "stdout": e.stdout,
            "stderr": e.stderr
        }

@mcp.tool()
def bactopia_tool(
    tool_name: str,
    bactopia_dir: str,
    outdir: str,
    profile: str = "conda",
    cpus: int = 2,
    extra_params: Optional[List[str]] = None
) -> dict:
    """
    Run a Bactopia Tool (comparative analysis) on existing Bactopia outputs.
    Common tools include: pangenome, roary, iqtree, pirate, summary, etc.
    
    Args:
        tool_name: Name of the Bactopia Tool to run.
        bactopia_dir: Directory containing previous Bactopia results.
        outdir: Directory to store tool results.
        profile: Nextflow profile to use.
        cpus: Number of CPUs to allocate.
        extra_params: List of additional command line arguments for the specific tool.
    """
    b_path = Path(bactopia_dir)
    o_path = Path(outdir)
    
    if not b_path.exists():
        return {"error": f"Bactopia results directory not found at {bactopia_dir}"}

    cmd = [
        "bactopia", tool_name,
        "--bactopia", str(b_path),
        "--outdir", str(o_path),
        "-profile", profile,
        "--max_cpus", str(cpus)
    ]
    
    if extra_params:
        cmd.extend(extra_params)

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        return {
            "command_executed": " ".join(cmd),
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_directory": str(o_path.absolute())
        }
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": " ".join(cmd),
            "error": str(e),
            "stdout": e.stdout,
            "stderr": e.stderr
        }

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