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

from mcp.server.fastmcp import FastMCP

SERVER_NAME = 'local_bin2cell'
mcp = FastMCP(SERVER_NAME)

@mcp.tool()
def bin2cell_prepare_bins(
    input_dir: str,
    output_h5ad: str,
    bin_size: int = 2,
    destripe: bool = True,
):
    """
    Reads Visium HD data from a SpaceRanger output directory and optionally performs destriping 
    to correct for technical effects in 2um bin data.

    Args:
        input_dir: Path to the SpaceRanger output directory (containing 'outs' folder).
        output_h5ad: Path where the processed bin-level AnnData object will be saved.
        bin_size: Resolution of bins to load (default is 2um).
        destripe: Whether to apply the destriping correction for variable bin dimensions.
    """
    # Input validation
    input_path = Path(input_dir)
    if not input_path.exists():
        return {"error": f"Input directory {input_dir} does not exist."}
    
    output_path = Path(output_h5ad)
    if not output_path.parent.exists():
        output_path.parent.mkdir(parents=True, exist_ok=True)

    if bin_size <= 0:
        return {"error": "bin_size must be a positive integer."}

    # Construct Python command
    # We use a python script string to execute the library functions
    destripe_cmd = "b2c.pp.destripe(adata)" if destripe else "pass"
    python_script = f"""
import bin2cell as b2c
import scanpy as sc
import os

try:
    # Load Visium HD data
    adata = b2c.pp.read_visium_hd_folder('{input_dir}', bin_size={bin_size})
    
    # Perform destriping if requested
    if {destripe}:
        b2c.pp.destripe(adata)
    
    # Save the result
    adata.write('{output_h5ad}')
    print("Successfully prepared bin data.")
except Exception as e:
    print(f"Error: {{str(e)}}")
    exit(1)
"""

    try:
        result = subprocess.run(
            ["python", "-c", python_script],
            capture_output=True,
            text=True,
            check=True
        )
        return {
            "command_executed": f"bin2cell.pp.read_visium_hd_folder and destripe={destripe}",
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": [output_h5ad]
        }
    except subprocess.CalledProcessError as e:
        return {
            "error": "Failed to prepare bin data",
            "stdout": e.stdout,
            "stderr": e.stderr,
            "command_executed": e.cmd
        }

@mcp.tool()
def bin2cell_run_stardist(
    image_path: str,
    output_mask_path: str,
    model_name: str = "2D_versatile_he",
    prob_thresh: float = 0.5,
    nms_thresh: float = 0.3,
):
    """
    Performs cell segmentation on a morphology image using StarDist.

    Args:
        image_path: Path to the high-resolution morphology image (e.g., tissue_hires_image.png).
        output_mask_path: Path where the resulting segmentation mask (.tif) will be saved.
        model_name: StarDist model to use (default: '2D_versatile_he').
        prob_thresh: Probability threshold for StarDist detection.
        nms_thresh: Non-maximum suppression threshold for StarDist.
    """
    # Input validation
    img_path = Path(image_path)
    if not img_path.exists():
        return {"error": f"Image file {image_path} does not exist."}
    
    out_mask = Path(output_mask_path)
    if not out_mask.parent.exists():
        out_mask.parent.mkdir(parents=True, exist_ok=True)

    python_script = f"""
import bin2cell as b2c
import cv2
import numpy as np

try:
    # Run StarDist segmentation via bin2cell wrapper
    # Note: bin2cell.tl.stardist handles the model loading and prediction
    b2c.tl.stardist(
        '{image_path}', 
        '{output_mask_path}', 
        model='{model_name}', 
        prob_thresh={prob_thresh}, 
        nms_thresh={nms_thresh}
    )
    print("Successfully generated segmentation mask.")
except Exception as e:
    print(f"Error: {{str(e)}}")
    exit(1)
"""

    try:
        result = subprocess.run(
            ["python", "-c", python_script],
            capture_output=True,
            text=True,
            check=True
        )
        return {
            "command_executed": f"bin2cell.tl.stardist on {image_path}",
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": [output_mask_path]
        }
    except subprocess.CalledProcessError as e:
        return {
            "error": "StarDist segmentation failed",
            "stdout": e.stdout,
            "stderr": e.stderr,
            "command_executed": e.cmd
        }

@mcp.tool()
def bin2cell_extract_cells(
    bin_h5ad_path: str,
    mask_path: str,
    output_cell_h5ad_path: str,
    qc_metrics: bool = True,
):
    """
    Groups subcellular bins into cells based on a segmentation mask and generates a cell-level AnnData object.

    Args:
        bin_h5ad_path: Path to the bin-level AnnData object (output from bin2cell_prepare_bins).
        mask_path: Path to the segmentation mask file (output from bin2cell_run_stardist).
        output_cell_h5ad_path: Path where the final cell-level AnnData object will be saved.
        qc_metrics: Whether to calculate standard scanpy QC metrics for the new cell object.
    """
    # Input validation
    bin_path = Path(bin_h5ad_path)
    if not bin_path.exists():
        return {"error": f"Bin h5ad file {bin_h5ad_path} does not exist."}
    
    m_path = Path(mask_path)
    if not m_path.exists():
        return {"error": f"Mask file {mask_path} does not exist."}

    out_cell_path = Path(output_cell_h5ad_path)
    if not out_cell_path.parent.exists():
        out_cell_path.parent.mkdir(parents=True, exist_ok=True)

    python_script = f"""
import bin2cell as b2c
import scanpy as sc

try:
    # Load the bin-level data
    adata_bins = sc.read_h5ad('{bin_h5ad_path}')
    
    # Extract cells based on the mask
    adata_cells = b2c.tl.extract_cells(adata_bins, '{mask_path}')
    
    # Calculate QC metrics if requested
    if {qc_metrics}:
        sc.pp.calculate_qc_metrics(adata_cells, inplace=True)
    
    # Save the cell-level object
    adata_cells.write('{output_cell_h5ad_path}')
    print("Successfully extracted cells from bins.")
except Exception as e:
    print(f"Error: {{str(e)}}")
    exit(1)
"""

    try:
        result = subprocess.run(
            ["python", "-c", python_script],
            capture_output=True,
            text=True,
            check=True
        )
        return {
            "command_executed": f"bin2cell.tl.extract_cells using mask {mask_path}",
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": [output_cell_h5ad_path]
        }
    except subprocess.CalledProcessError as e:
        return {
            "error": "Cell extraction failed",
            "stdout": e.stdout,
            "stderr": e.stderr,
            "command_executed": e.cmd
        }

@mcp.tool()
def bin2cell_visualize_segmentation(
    bin_h5ad_path: str,
    output_image_path: str,
    basis: str = "spatial",
):
    """
    Generates a visualization of the bin-to-cell assignments.

    Args:
        bin_h5ad_path: Path to the bin-level AnnData object containing cell assignments.
        output_image_path: Path to save the visualization plot (e.g., .png or .pdf).
        basis: The coordinate system to use for plotting (default: 'spatial').
    """
    bin_path = Path(bin_h5ad_path)
    if not bin_path.exists():
        return {"error": f"Bin h5ad file {bin_h5ad_path} does not exist."}

    python_script = f"""
import scanpy as sc
import matplotlib.pyplot as plt
import bin2cell as b2c

try:
    adata = sc.read_h5ad('{bin_h5ad_path}')
    if 'cell_id' not in adata.obs.columns:
        print("Error: cell_id not found in adata.obs. Run extract_cells first.")
        exit(1)
        
    # Plotting logic
    sc.pl.embedding(adata, basis='{basis}', color='cell_id', show=False)
    plt.savefig('{output_image_path}')
    print("Successfully saved visualization.")
except Exception as e:
    print(f"Error: {{str(e)}}")
    exit(1)
"""

    try:
        result = subprocess.run(
            ["python", "-c", python_script],
            capture_output=True,
            text=True,
            check=True
        )
        return {
            "command_executed": f"Visualization of cell_id on {basis}",
            "stdout": result.stdout,
            "stderr": result.stderr,
            "output_files": [output_image_path]
        }
    except subprocess.CalledProcessError as e:
        return {
            "error": "Visualization failed",
            "stdout": e.stdout,
            "stderr": e.stderr,
            "command_executed": e.cmd
        }

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