import os
import argparse
import subprocess
import json
from pathlib import Path
from typing import Dict, List, Optional

from bioinfomcp_converter import BioinfoMCP


def generate_requirements_with_pipreqs(tool_name, server_path):
    """
    Use pipreqs to generate requirements.txt (fallback to minimal deps when pipreqs fails).
    """
    _ = tool_name
    app_path = Path(server_path) / "app"
    target_path = app_path if app_path.exists() else Path(server_path)
    try:
        subprocess.run(
            ["pipreqs", str(target_path), "--force"],
            capture_output=True,
            text=True,
            check=True,
        )
        print(f"Generated requirements.txt using pipreqs from: {target_path}")
    except FileNotFoundError:
        req_path = Path(server_path) / "requirements.txt"
        req_path.write_text("fastmcp\n", encoding="utf-8")
        print("pipreqs not found, generated fallback requirements.txt")
    except subprocess.CalledProcessError as exc:
        req_path = Path(server_path) / "requirements.txt"
        req_path.write_text("fastmcp\n", encoding="utf-8")
        stderr = (exc.stderr or "").strip()
        stdout = (exc.stdout or "").strip()
        detail = stderr or stdout or "no stderr/stdout"
        print(f"pipreqs failed ({detail}), generated fallback requirements.txt")
    except Exception as exc:
        req_path = Path(server_path) / "requirements.txt"
        req_path.write_text("fastmcp\n", encoding="utf-8")
        print(f"pipreqs failed ({exc}), generated fallback requirements.txt")
    return Path(server_path) / "requirements.txt"

def generate_environment_yaml(tool_name, server_path):
    """
    generate environment.yaml
    """
    result = f"""
name: mcp-tool
channels:
  - bioconda
  - conda-forge
  - defaults
dependencies:
  - {tool_name}
  - python=3.10
    """
    with open(server_path / "environment.yaml", "w") as f:
        f.write(result)
    print("Generated environment.yaml")
    return Path(server_path) / "environment.yaml"

def generate_environment_yml(tool_name, server_path):
    """
    Generate requirements.yml
    """
    _ = server_path
    result = f"""name: {tool_name}_env
channels:
  - bioconda
  - conda-forge
dependencies:
  - python=3.10
  - {tool_name}
  - pip
"""
    return result

def parse_bool(value) -> bool:
    if isinstance(value, bool):
        return value
    if value is None:
        return False
    return str(value).strip().lower() in ("1", "true", "yes", "y", "on")


def has_existing_server(output_root: Path, tool_name: str) -> Optional[Path]:
    server_file = output_root / f"mcp_{tool_name}" / "app" / f"{tool_name}_server.py"
    if server_file.exists() and server_file.is_file() and server_file.stat().st_size > 0:
        return server_file
    return None


def convert_mcptool(tool_name, manual, run_help_command, server_path, converter=None):
    converter = converter or BioinfoMCP()
    app_path = Path(server_path) / "app"
    app_path.mkdir(parents=True, exist_ok=True)

    conversion_log: Dict = {
        "tool_name": tool_name,
        "manual": manual,
        "run_help_command": bool(run_help_command),
        "status": "failed",
        "error": "",
        "output_file": "",
    }

    try:
        conv_result = converter.autogenerate_mcp_tool(tool_name, manual, run_help_command)
    except Exception as exc:
        conversion_log["error"] = f"autogenerate_mcp_tool failed: {exc}"
        return conversion_log
    retry = 0
    while not conv_result[0] and retry < 6:
        retry += 1
        if conv_result[2] is None:
            try:
                conv_result = converter.autogenerate_mcp_tool(tool_name, manual, run_help_command)
                print(f"{tool_name}: regenerate ({retry})")
            except Exception as exc:
                conversion_log["error"] = f"regenerate failed (retry={retry}): {exc}"
                return conversion_log
        else:
            try:
                conv_result = converter.refine_after_feedback(tool_name, code=conv_result[2], error_message=conv_result[1])
                print(f"{tool_name}: refine ({retry})")
            except Exception as exc:
                conversion_log["error"] = f"refine failed (retry={retry}): {exc}"
                return conversion_log

    if not conv_result[0] or conv_result[2] is None:
        conversion_log["error"] = conv_result[1] or "Unknown conversion failure"
        return conversion_log

    output_path = Path(server_path) / "app" / f"{tool_name}_server.py"
    output_path.write_text(conv_result[2], encoding="utf-8")
    conversion_log["status"] = "success"
    conversion_log["output_file"] = str(output_path)
    return conversion_log


def load_tool_list_from_json(json_path: Path, top_k: int = 0, domains: Optional[List[str]] = None) -> List[str]:
    with open(json_path, "r", encoding="utf-8") as f:
        rows = json.load(f)
    if not isinstance(rows, list):
        raise ValueError(f"Expected list JSON: {json_path}")

    domains = [d.strip().lower() for d in (domains or []) if d.strip()]
    tools: List[str] = []
    for row in rows:
        if not isinstance(row, dict):
            continue
        if domains:
            domain = str(row.get("domain", "")).strip().lower()
            if domain not in domains:
                continue
        package_name = str(row.get("package_name", "")).strip()
        software_name = str(row.get("software_name", "")).strip()
        tool_name = package_name or software_name
        if tool_name:
            tools.append(tool_name)

    deduped = list(dict.fromkeys(tools))
    if top_k > 0:
        deduped = deduped[:top_k]
    return deduped


def load_jobs_from_json(json_path: Path, top_k: int = 0, domains: Optional[List[str]] = None) -> List[Dict]:
    with open(json_path, "r", encoding="utf-8") as f:
        rows = json.load(f)
    if not isinstance(rows, list):
        raise ValueError(f"Expected list JSON: {json_path}")

    domains = [d.strip().lower() for d in (domains or []) if d.strip()]
    jobs: List[Dict] = []
    for row in rows:
        if not isinstance(row, dict):
            continue
        name = (
            str(row.get("name", "")).strip()
            or str(row.get("package_name", "")).strip()
            or str(row.get("software_name", "")).strip()
        )
        if not name:
            continue
        domain = str(row.get("domain", "")).strip().lower()
        if domains and domain not in domains:
            continue
        jobs.append(
            {
                "name": name,
                "manual": str(row.get("manual", "--help")),
                "run_help_command": parse_bool(row.get("run_help_command", True)),
                "domain": domain,
                "tier": str(row.get("tier", "")),
            }
        )
    if top_k > 0:
        jobs = jobs[:top_k]
    return jobs


def load_jobs_from_manual_dir(manual_dir: Path, top_k: int = 0) -> List[Dict]:
    """
    Load conversion jobs from a directory containing *.help.txt or *.manual_bundle.txt.
    """
    if not manual_dir.exists():
        raise FileNotFoundError(f"Manual directory not found: {manual_dir}")

    candidates = sorted(
        list(manual_dir.glob("*.help.txt")) + list(manual_dir.glob("*.manual_bundle.txt"))
    )
    # De-duplicate while preferring manual_bundle when both exist.
    by_tool: Dict[str, Path] = {}
    for p in candidates:
        name = p.name
        if name.endswith(".manual_bundle.txt"):
            tool_name = name[: -len(".manual_bundle.txt")]
            by_tool[tool_name] = p
        elif name.endswith(".help.txt"):
            tool_name = name[: -len(".help.txt")]
            by_tool.setdefault(tool_name, p)

    jobs: List[Dict] = []
    for tool_name, manual_path in sorted(by_tool.items()):
        jobs.append(
            {
                "name": tool_name,
                "manual": str(manual_path),
                "run_help_command": False,
                "domain": "",
                "tier": "",
            }
        )
    if top_k > 0:
        jobs = jobs[:top_k]
    return jobs

def select_jobs_for_shard(jobs: List[Dict], shard_count: int, shard_index: int) -> List[Dict]:
    """
    Deterministically split jobs for multi-terminal parallel execution.
    """
    if shard_count <= 1:
        return jobs
    if shard_index < 0 or shard_index >= shard_count:
        raise ValueError(f"Invalid shard_index={shard_index}. Expected 0 <= shard_index < {shard_count}")
    selected: List[Dict] = []
    for idx, job in enumerate(jobs):
        if idx % shard_count == shard_index:
            selected.append(job)
    return selected

def dockerfile_content(tool_name):
    return f"""
FROM python:3.10-slim

# Install system dependencies
RUN apt-get update && \
    apt-get install -y \
        default-jre \
        wget \
        curl \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

# Install Miniconda
RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/miniconda.sh && \
    bash /tmp/miniconda.sh -b -p /opt/conda && \
    rm /tmp/miniconda.sh

# Add conda to PATH
ENV PATH="/opt/conda/bin:$PATH"

# Install {tool_name} via conda (e.g., from bioconda)
RUN conda install -c bioconda {tool_name} -y && \
    conda clean -a

# Install Python dependencies
RUN pip install uv
RUN uv pip install --system fastmcp

# Create app directory
WORKDIR /app

# Copy your MCP server
COPY app/{tool_name}_server.py /app/

# Create workspace and output directories
RUN mkdir -p /app/workspace /app/output

# Make sure the server script is executable
RUN chmod +x /app/{tool_name}_server.py

# Expose port for MCP over HTTP (optional)
EXPOSE 8000

# Health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
    CMD python -c "import sys; sys.exit(0)"

# Default command runs the MCP server via stdio
CMD ["python", "/app/{tool_name}_server.py"]
        """

def dockercompose_content(tool_name):
    """Generate docker-compose.yml for easy deployment"""
    
    compose = f"""version: '3.8'

services:
  mcp-{tool_name}:
    build: .
    image: mcp-{tool_name}:latest
    container_name: mcp-{tool_name}
    ports:
      - "8000:8000"
    environment:
      - MCP_SERVER_NAME={tool_name}
    volumes:
      - ./workspace:/app/workspace
      - ./output:/app/output
    restart: unless-stopped
    healthcheck:
      test: ["CMD", "python", "-c", "import sys; sys.exit(0)"]
      interval: 30s
      timeout: 10s
      retries: 3
      start_period: 5s
    """

    return compose

def image_name_for_tool(tool_name: str) -> str:
    return f"mcp-{tool_name}:latest"


def build_docker_image(tool_name, server_path, output_path, is_pipeline):

    # Create build directory
    try:
        _ = output_path
        df_content = dockerfile_content(tool_name)
        with open(server_path / "Dockerfile", "w") as f:
            f.write(df_content) 
                
         # make the docker-compose.yml
        if not is_pipeline:
            dc_content = dockercompose_content(tool_name)
            with open(server_path/"docker-compose.yml", "w") as f:
                f.write(dc_content)
        
        # subprocess.run(["docker", "build", "-t", f"{args.name}-docker", server_path])

        return 1
    except:
        return 0


def claude_addition(tool_name, server_path):
    image_name = image_name_for_tool(tool_name)
    workspace_path = str((Path(server_path) / "workspace").resolve())
    config = {
        "mcpServers": {
            tool_name: {
                "command": "docker",
                "args": [
                    "run",
                    "--rm",
                    "-i",
                    "-v", 
                    f"{workspace_path}:/app/workspace",
                    image_name
                ]
            }
        }
    }
    return json.dumps(config, indent=2)


if __name__ == '__main__':
    parser = argparse.ArgumentParser("Accept the Bioinformatic tool name and the help document")
    parser.add_argument('--model', type=str, default='openai', choices=['openai', 'azure', 'gemini'], help="LLM backend")
    parser.add_argument('--name', type=str)
    parser.add_argument('--manual', type=str, default='--help', help="The file path to the help document or help flag")
    parser.add_argument('--run_help_command', type=str, default='True')
    parser.add_argument('--output_location', type=str, required=True)
    parser.add_argument('--is_pipeline', action='store_true', default=False)
    parser.add_argument('--tools_json', type=str, default="", help="JSON list file from crawler output")
    parser.add_argument('--jobs_json', type=str, default="", help="Per-tool jobs JSON with name/manual/run_help_command")
    parser.add_argument('--manual_dir', type=str, default="", help="Directory containing *.help.txt or *.manual_bundle.txt")
    parser.add_argument('--top_k', type=int, default=0, help="Only convert first K tools from tools_json/jobs_json")
    parser.add_argument('--domains', nargs='*', default=[], help="Domain filter for tools_json/jobs_json")
    parser.add_argument('--skip_existing', type=str, default='True', help="Skip tool if MCP server file already exists")
    parser.add_argument('--shard_count', type=int, default=1, help="Total shards for multi-terminal execution")
    parser.add_argument('--shard_index', type=int, default=0, help="Current shard index in [0, shard_count)")
    parser.add_argument('--report_suffix', type=str, default="", help="Optional suffix for report filename")
    args = parser.parse_args()

    output_root = Path(args.output_location)
    output_root.mkdir(parents=True, exist_ok=True)
    run_help_command = parse_bool(args.run_help_command)
    skip_existing = parse_bool(args.skip_existing)
    converter = BioinfoMCP(model=args.model)

    jobs: List[Dict] = []
    if args.jobs_json:
        jobs = load_jobs_from_json(Path(args.jobs_json), top_k=args.top_k, domains=args.domains)
    elif args.manual_dir:
        jobs = load_jobs_from_manual_dir(Path(args.manual_dir), top_k=args.top_k)
    elif args.tools_json:
        tools = load_tool_list_from_json(Path(args.tools_json), top_k=args.top_k, domains=args.domains)
        jobs = [{"name": t, "manual": args.manual, "run_help_command": run_help_command} for t in tools]
    elif args.name:
        jobs = [{"name": args.name, "manual": args.manual, "run_help_command": run_help_command}]
    else:
        raise ValueError("Either --name, --tools_json, --jobs_json or --manual_dir must be provided")

    if args.shard_count < 1:
        raise ValueError(f"--shard_count must be >= 1, got: {args.shard_count}")
    jobs = select_jobs_for_shard(jobs, args.shard_count, args.shard_index)
    print(f"Shard selection: shard_index={args.shard_index}/{args.shard_count}, jobs={len(jobs)}")

    run_log: List[Dict] = []
    for job in jobs:
        tool_name = job["name"]
        manual = job.get("manual", "--help")
        job_run_help = parse_bool(job.get("run_help_command", True))

        existing_server = has_existing_server(output_root, tool_name) if skip_existing else None
        if existing_server:
            print(f"Skip existing MCP server: {existing_server}")
            run_log.append(
                {
                    "tool_name": tool_name,
                    "manual": manual,
                    "run_help_command": bool(job_run_help),
                    "status": "skipped_existing",
                    "error": "",
                    "output_file": str(existing_server),
                    "dockerfile_ready": True,
                }
            )
            continue

        server_path = output_root / f"mcp_{tool_name}"
        os.makedirs(server_path, exist_ok=True)
        print("Server path", server_path)
        app_path = Path(server_path) / "app"
        os.makedirs(app_path, exist_ok=True)

        result = convert_mcptool(tool_name, manual, job_run_help, server_path, converter=converter)
        if result.get("status") == "success":
            generate_requirements_with_pipreqs(tool_name, server_path)
            generate_environment_yaml(tool_name, server_path)

            status = build_docker_image(tool_name, server_path, output_root, args.is_pipeline)
            result["dockerfile_ready"] = bool(status)

            if status:
                add = claude_addition(tool_name, server_path)
                print(f"Add the following onto your Claude Configuration json file to run the MCP server:\n{'=='*10}\n{add}\n{'=='*10}")
            else:
                print("Failed to build Docker Image")
        else:
            result["dockerfile_ready"] = False
            print(f"Skip packaging for {tool_name}: conversion failed -> {result.get('error', '')}")

        run_log.append(result)

    report_name = "conversion_report.json"
    if args.shard_count > 1:
        report_name = f"conversion_report.shard_{args.shard_index}_of_{args.shard_count}.json"
    if args.report_suffix:
        report_name = report_name.replace(".json", f".{args.report_suffix}.json")
    report_path = output_root / report_name
    report_path.write_text(json.dumps(run_log, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"Saved conversion report to: {report_path}")