import subprocess
import os
import re
from pathlib import Path
from mcp.server.fastmcp import FastMCP

# ── Initialize MCP Server ─────────────────────────────────────────────────────
mcp = FastMCP(
    name="speckit-plus"
)

# ── Security: Allowed AI agents whitelist ─────────────────────────────────────
ALLOWED_AI_AGENTS = frozenset([
    "qwen", "claude", "gemini", "copilot", 
    "cursor", "windsurf", "opencode", "codex", "amp"
])

# ── Security: Maximum input lengths ───────────────────────────────────────────
MAX_PATH_LENGTH = 255
MAX_CONTENT_LENGTH = 10000  # Max characters for principles/requirements/tech_stack

# ── Security: Path traversal pattern detection ────────────────────────────────
# Matches: ..  or ../ or ..\  but NOT single dots
PATH_TRAVERSAL_PATTERN = re.compile(r'\.\.[\\/]|^[.]{2}$|[\\/]\.\.[\\/]|[\\/]\.\.$')


def _validate_project_path(project_path: str) -> tuple[bool, str]:
    """
    Validate project path to prevent path traversal attacks.
    
    Returns:
        Tuple of (is_valid, error_message)
    """
    if not project_path:
        return False, "Project path cannot be empty"
    
    if len(project_path) > MAX_PATH_LENGTH:
        return False, f"Project path exceeds maximum length of {MAX_PATH_LENGTH}"
    
    # Check for path traversal patterns
    if PATH_TRAVERSAL_PATTERN.search(project_path):
        return False, "Path traversal patterns (..) are not allowed"
    
    # Check for null bytes
    if '\x00' in project_path:
        return False, "Null bytes are not allowed in path"
    
    # Resolve to absolute path and validate
    try:
        resolved_path = Path(project_path).resolve()
        
        # Ensure path exists
        if not resolved_path.exists():
            return False, f"Path does not exist: {project_path}"
        
        # Ensure it's a directory
        if not resolved_path.is_dir():
            return False, f"Path is not a directory: {project_path}"
        
        # Check for .specify folder (indicates initialized Spec-Kit project)
        specify_folder = resolved_path / ".specify"
        if not specify_folder.exists():
            return False, f"Not an initialized Spec-Kit project: {project_path}. Run 'sp init' first."
        
    except (OSError, ValueError) as e:
        return False, f"Invalid path: {str(e)}"
    
    return True, str(resolved_path)


def _validate_ai_agent(ai_agent: str) -> tuple[bool, str]:
    """
    Validate AI agent name against whitelist.
    
    Returns:
        Tuple of (is_valid, error_message)
    """
    if not ai_agent:
        return False, "AI agent cannot be empty"
    
    if ai_agent.lower() not in ALLOWED_AI_AGENTS:
        return False, f"Invalid AI agent. Allowed: {', '.join(sorted(ALLOWED_AI_AGENTS))}"
    
    return True, ai_agent.lower()


def _validate_content(content: str, field_name: str) -> tuple[bool, str]:
    """
    Validate content for file writing operations.
    
    Returns:
        Tuple of (is_valid, error_message)
    """
    if not content:
        return False, f"{field_name} cannot be empty"
    
    if len(content) > MAX_CONTENT_LENGTH:
        return False, f"{field_name} exceeds maximum length of {MAX_CONTENT_LENGTH} characters"
    
    # Check for null bytes
    if '\x00' in content:
        return False, f"Null bytes are not allowed in {field_name}"
    
    return True, content


def _safe_join_path(base_path: str, *subpaths: str) -> str:
    """
    Safely join path components and ensure result stays within base path.
    
    Returns:
        Resolved absolute path string
    """
    base = Path(base_path).resolve()
    result = base
    
    for subpath in subpaths:
        # Sanitize subpath - remove any traversal attempts
        clean_subpath = re.sub(r'\.\.[\\/]', '', subpath)
        result = result / clean_subpath
    
    # Ensure final path is still within base
    try:
        result.resolve().relative_to(base)
    except ValueError:
        raise ValueError(f"Path traversal detected: {subpaths}")
    
    return str(result)


# ── Helper: run shell commands ────────────────────────────────────────────────
def run_sp(args: list[str], cwd: str = None) -> str:
    """
    Run a specifyplus (sp) command and return output.
    
    Security: Uses list form (NOT shell=True) to prevent shell injection.
    """
    # Validate args don't contain dangerous patterns
    for arg in args:
        if '\x00' in arg:
            return "❌ Error: Null bytes not allowed in command arguments"
    
    try:
        result = subprocess.run(
            ["sp"] + args,
            capture_output=True,
            text=True,
            cwd=cwd or os.getcwd(),
            timeout=300  # 5 minute timeout
        )
        if result.returncode == 0:
            return result.stdout.strip() or "✅ Command completed successfully."
        else:
            return f"❌ Error:\n{result.stderr.strip()}"
    except subprocess.TimeoutExpired:
        return "❌ Error: Command timed out after 5 minutes"
    except FileNotFoundError:
        return "❌ 'sp' command not found. Run: pip install specifyplus"
    except Exception as e:
        # Sanitize error message to avoid leaking sensitive info
        error_msg = str(e).replace(os.getcwd(), "<cwd>")
        return f"❌ Error: {type(e).__name__} occurred"


# ── TOOL 1: Initialize a new project ─────────────────────────────────────────
@mcp.tool()
def speckit_init(
    project_name: str,
    ai_agent: str = "qwen",
    current_directory: bool = False
) -> str:
    """
    Initialize a new Spec-Kit Plus project.

    Args:
        project_name: Name of the new project folder (use '.' for current dir)
        ai_agent: AI agent to configure for. Options: qwen, claude, gemini,
                  copilot, cursor, windsurf, opencode, codex, amp
        current_directory: If True, initialize in current directory
    """
    # Security: Validate AI agent against whitelist
    is_valid, result = _validate_ai_agent(ai_agent)
    if not is_valid:
        return f"❌ Invalid AI agent: {result}"
    validated_agent = result
    
    # Security: Validate project_name length and patterns
    if not project_name or len(project_name) > MAX_PATH_LENGTH:
        return f"❌ Invalid project name: exceeds maximum length"
    if PATH_TRAVERSAL_PATTERN.search(project_name):
        return "❌ Invalid project name: contains path traversal patterns"
    if '\x00' in project_name:
        return "❌ Invalid project name: contains null bytes"
    
    args = ["init"]

    if current_directory:
        args += ["--here", "--force", "--ai", validated_agent]
    else:
        args += [project_name, "--ai", validated_agent]

    return run_sp(args)


# ── TOOL 2: Check system requirements ────────────────────────────────────────
@mcp.tool()
def speckit_check() -> str:
    """
    Check that all required tools are installed (git, sp, AI agent CLIs).
    Run this first to verify your setup is working.
    """
    return run_sp(["check"])


# ── TOOL 3: Create project constitution ──────────────────────────────────────
@mcp.tool()
def speckit_constitution(
    principles: str,
    project_path: str = "."
) -> str:
    """
    Create or update the project's governing principles (constitution.md).
    This guides all AI decisions during development.

    Args:
        principles: Description of the principles you want to enforce.
                    Example: "Focus on clean code, testing, and performance"
        project_path: Path to the initialized project folder
    """
    # Security: Validate project path
    is_valid, result = _validate_project_path(project_path)
    if not is_valid:
        return f"❌ Invalid project path: {result}"
    validated_path = result
    
    # Security: Validate principles content
    is_valid, result = _validate_content(principles, "Principles")
    if not is_valid:
        return f"❌ Invalid principles: {result}"
    validated_principles = result
    
    # Security: Use safe path joining
    try:
        specify_folder = _safe_join_path(validated_path, ".specify")
        prompt_file = os.path.join(specify_folder, "constitution_prompt.txt")
        
        os.makedirs(specify_folder, exist_ok=True)
        with open(prompt_file, "w", encoding="utf-8") as f:
            f.write(validated_principles)
    except (ValueError, OSError) as e:
        return f"❌ Error writing constitution file: {type(e).__name__}"

    return (
        f"✅ Constitution principles saved to: {prompt_file}\n\n"
        f"📋 Next step: Open your AI agent in '{validated_path}' and run:\n"
        f"   /sp.constitution {principles[:100]}..."
    )


# ── TOOL 4: Write specification (WHAT to build) ───────────────────────────────
@mcp.tool()
def speckit_specify(
    requirements: str,
    project_path: str = "."
) -> str:
    """
    Write a project specification — describe WHAT you want to build and WHY.
    Do NOT mention tech stack here. Focus on features and user stories.

    Args:
        requirements: Plain English description of what to build.
                      Example: "Build a task manager with Kanban boards,
                      user assignments, and drag-and-drop."
        project_path: Path to the initialized project folder
    """
    # Security: Validate project path
    is_valid, result = _validate_project_path(project_path)
    if not is_valid:
        return f"❌ Invalid project path: {result}"
    validated_path = result
    
    # Security: Validate requirements content
    is_valid, result = _validate_content(requirements, "Requirements")
    if not is_valid:
        return f"❌ Invalid requirements: {result}"
    validated_requirements = result
    
    # Security: Use safe path joining
    try:
        specify_folder = _safe_join_path(validated_path, ".specify")
        spec_prompt_path = os.path.join(specify_folder, "spec_input.txt")
        
        os.makedirs(specify_folder, exist_ok=True)
        with open(spec_prompt_path, "w", encoding="utf-8") as f:
            f.write(validated_requirements)
    except (ValueError, OSError) as e:
        return f"❌ Error writing specification file: {type(e).__name__}"

    return (
        f"✅ Specification input saved.\n\n"
        f"📋 Next step: Open your AI agent in '{validated_path}' and run:\n"
        f"   /sp.specify {requirements[:100]}..."
    )


# ── TOOL 5: Generate technical plan (HOW to build it) ─────────────────────────
@mcp.tool()
def speckit_plan(
    tech_stack: str,
    project_path: str = "."
) -> str:
    """
    Generate a technical implementation plan from the specification.
    This is where you specify your tech stack and architecture.

    Args:
        tech_stack: Your chosen technologies and architecture.
                    Example: "FastAPI backend, PostgreSQL, React frontend,
                    Docker deployment"
        project_path: Path to the initialized project folder
    """
    # Security: Validate project path
    is_valid, result = _validate_project_path(project_path)
    if not is_valid:
        return f"❌ Invalid project path: {result}"
    validated_path = result
    
    # Security: Validate tech_stack content
    is_valid, result = _validate_content(tech_stack, "Tech stack")
    if not is_valid:
        return f"❌ Invalid tech stack: {result}"
    validated_tech_stack = result
    
    # Security: Use safe path joining
    try:
        specify_folder = _safe_join_path(validated_path, ".specify")
        plan_input_path = os.path.join(specify_folder, "plan_input.txt")
        
        os.makedirs(specify_folder, exist_ok=True)
        with open(plan_input_path, "w", encoding="utf-8") as f:
            f.write(validated_tech_stack)
    except (ValueError, OSError) as e:
        return f"❌ Error writing plan file: {type(e).__name__}"

    return (
        f"✅ Tech stack input saved.\n\n"
        f"📋 Next step: Open your AI agent in '{validated_path}' and run:\n"
        f"   /sp.plan {tech_stack[:100]}..."
    )


# ── TOOL 6: Break Plan into tasks ────────────────────────────────────────────
@mcp.tool()
def speckit_tasks(project_path: str = ".") -> str:
    """
    Break the technical implementation plan into an ordered task list.
    Run this AFTER sp.plan is complete.

    Args:
        project_path: Path to the initialized project folder
    """
    # Security: Validate project path
    is_valid, result = _validate_project_path(project_path)
    if not is_valid:
        return f"❌ Invalid project path: {result}"
    validated_path = result
    
    return (
        f"📋 Next step: Open your AI agent in '{validated_path}' and run:\n"
        f"   /sp.tasks\n\n"
        f"This will generate a tasks.md file with all implementation steps."
    )


# ── TOOL 7: Implement the project ─────────────────────────────────────────────
@mcp.tool()
def speckit_implement(project_path: str = ".") -> str:
    """
    Execute all tasks from the task list to build the full project.
    Run this AFTER /sp.tasks is complete.

    Args:
        project_path: Path to the initialized project folder
    """
    # Security: Validate project path
    is_valid, result = _validate_project_path(project_path)
    if not is_valid:
        return f"❌ Invalid project path: {result}"
    validated_path = result
    
    return (
        f"📋 Next step: Open your AI agent in '{validated_path}' and run:\n"
        f"   /sp.implement\n\n"
        f"The AI will execute all tasks in order and build your project."
    )


# ── TOOL 8: Clarify requirements ──────────────────────────────────────────────
@mcp.tool()
def speckit_clarify(project_path: str = ".") -> str:
    """
    Run structured clarification BEFORE planning.
    The AI will ask questions to fill gaps in the specification.

    Args:
        project_path: Path to the initialized project folder
    """
    # Security: Validate project path
    is_valid, result = _validate_project_path(project_path)
    if not is_valid:
        return f"❌ Invalid project path: {result}"
    validated_path = result
    
    return (
        f"📋 Run in your AI agent inside '{validated_path}':\n"
        f"   /sp.clarify\n\n"
        f"The AI will ask targeted questions to improve your spec."
    )


# ── TOOL 9: Get full workflow guide ───────────────────────────────────────────
@mcp.tool()
def speckit_workflow_guide() -> str:
    """
    Get the complete Spec-Driven Development workflow steps in order.
    Use this to understand what commands to run and when.
    """
    return """
📘 SPEC-KIT PLUS — COMPLETE WORKFLOW GUIDE
==========================================

STEP 0: Initialize project
  Tool: speckit_init(project_name="my-app", ai_agent="qwen")
  CLI:  sp init my-app --ai qwen

STEP 1: Set principles (run inside AI agent)
  /sp.constitution  Enforce clean code, testing, performance standards

STEP 2: Write specification (WHAT to build, no tech stack)
  /sp.specify  Build a task manager with Kanban boards and user assignments

STEP 3: Clarify gaps (recommended before planning)
  /sp.clarify

STEP 4: Technical plan (HOW to build it, with tech stack)
  /sp.plan  Use FastAPI, PostgreSQL, React, Docker

STEP 5: Generate task list
  /sp.tasks

STEP 6: Build the project
  /sp.implement

SUPPORTED AI AGENTS:
  qwen, claude, gemini, copilot, cursor, windsurf, opencode, codex, amp

OUTPUT FILES (in .specify/):
  memory/constitution.md   ← Your principles
  specs/001-feature/
    spec.md                ← Requirements & user stories
    plan.md                ← Technical plan
    tasks.md               ← Task list
    data-model.md          ← Database schema
"""


# ── Run Server ────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    mcp.run()
