from mcp.server.fastmcp import FastMCP
import os
import json
import asyncio
import sys
import logging
import shutil
from pathlib import Path

# Additional platform context (used for path resolution reporting)
import platform

# ---------------------------------------------------------------------------
# Protocol Compliance Notes (JSON-RPC over newline-delimited stdio)
# ---------------------------------------------------------------------------
# This server communicates strictly via newline-delimited JSON-RPC 2.0 messages
# on stdout and expects the same on stdin. Each outbound message produced by
# the underlying FastMCP library is serialized to a single line with no
# embedded newline characters. Tool return values may contain newlines; these
# are escaped (e.g. "\n") inside the JSON string so the overall message line
# remains single-line compliant.
#
# The server:
#   * MUST NOT emit any non-JSON output on stdout. (All logging is routed to
#     stderr via Python's logging module.)
#   * MAY emit UTF-8 log lines to stderr. These are optional for clients.
#   * Accepts only well-formed JSON-RPC messages (one line per message). The
#     FastMCP library performs core validation; additional defensive logging
#     is added here for clarity.
# ---------------------------------------------------------------------------

# Configure logging to stderr explicitly to avoid accidental stdout writes.
logging.basicConfig(stream=sys.stderr, level=logging.INFO, format="[%(asctime)s] %(levelname)s: %(message)s")
logger = logging.getLogger(__name__)

# instantiate an MCP server for UBE build system
mcp = FastMCP("UBE Build System Server")

# Configuration
# We now support cross-platform resolution of the UBE executable. Resolution order:
#   1. Explicit UBE_PATH environment variable (always wins)
#   2. Platform-specific default installation locations
#   3. Fallback to PATH discovery via shutil.which
# For Linux we assume (customizable) default under $HOME/LocalEnvs/UBE/bin/ube or /home/wsluser/LocalEnvs/UBE/bin/ube.

def _resolve_ube_path():
    """Resolve the UBE executable path in a cross-platform manner.

    Returns:
        (path, source) where source is one of:
            'env' - provided via UBE_PATH environment variable
            'default-windows' / 'default-linux' - platform default found and exists
            'which' - discovered via PATH lookup
            'default-windows-missing' / 'default-linux-missing' - fallback path returned but not found
    """
    env_path = os.environ.get("UBE_PATH")
    if env_path:
        return env_path, "env"

    if os.name == "nt":  # Windows
        default_win = r"C:\Conan\aurix_rc1_sw\UBE\1.0.0\ube.exe"
        if os.path.isfile(default_win):
            return default_win, "default-windows"
        which_path = shutil.which("ube.exe") or shutil.which("ube")
        if which_path:
            return which_path, "which"
        return default_win, "default-windows-missing"

    # Linux / POSIX
    home = os.environ.get("HOME", "/home/wsluser")
    candidates = [
        f"{home}/LocalEnvs/UBE/bin/ube",
        "/home/wsluser/LocalEnvs/UBE/bin/ube",  # explicit WSL style fallback
        shutil.which("ube"),
    ]
    for c in candidates:
        if c and os.path.isfile(c):
            if "LocalEnvs" in c:
                return c, "default-linux"
            return c, "which"
    # None found; return first candidate as a missing default hint
    return candidates[0], "default-linux-missing"


UBE_PATH, UBE_PATH_SOURCE = _resolve_ube_path()
DEFAULT_WORKSPACE = os.environ.get("WORKSPACE_FOLDER", os.getcwd())

def _ube_exists(path: str) -> bool:
    """Check if path points to an executable file."""
    return os.path.isfile(path) and os.access(path, os.X_OK)

@mcp.tool()
async def ube_check(timeout_seconds: int = 5) -> str:
    """Check if the UBE toolchain is active by invoking a lightweight help command.

    Runs: ube.exe -h (help output)

    Parameters:
        timeout_seconds: (Deprecated / ignored) retained for backward compatibility.

    Returns:
        A short diagnostic string including exit code, first lines of STDOUT/STDERR
        or an actionable error message if the executable is missing/unresponsive.
    """
    if not _ube_exists(UBE_PATH):
        return (
            f"UBE executable not found or not executable at '{UBE_PATH}' (source={UBE_PATH_SOURCE}). "
            "Set UBE_PATH to a valid path."
        )
    try:
        # Run with just '-h' to avoid invalid subcommand errors; timeout not needed for quick help.
        args = ["-h"]
        logger.info("Starting UBE active check (no-timeout): %s %s", UBE_PATH, " ".join(args))
        process = await asyncio.create_subprocess_exec(
            UBE_PATH, *args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            stdin=asyncio.subprocess.DEVNULL,
        )
        stdout, stderr = await process.communicate()

        exit_code = process.returncode
        out_text = stdout.decode('utf-8', 'replace')
        err_text = stderr.decode('utf-8', 'replace')

        # Produce concise snippet (first 15 lines combined)
        def first_lines(text: str, limit: int = 15) -> str:
            lines = text.splitlines()
            snippet = lines[:limit]
            more = f"\n... ({len(lines)-limit} more lines)" if len(lines) > limit else ""
            return "\n".join(snippet) + more

        summary = f"Exit code: {exit_code}\n"
        if out_text:
            summary += "STDOUT (truncated):\n" + first_lines(out_text) + "\n"
        if err_text:
            summary += "STDERR (truncated):\n" + first_lines(err_text) + "\n"
        summary += f"ube_path={UBE_PATH} source={UBE_PATH_SOURCE}"
        return summary
    except Exception as e:
        logger.exception("UBE active check failed: %s", e)
        return f"Error invoking UBE help command: {e}"

# ---------------------------------------------------------------------------
# Long-running job management
# We provide non-blocking variants for lengthy UBE operations to avoid client
# request timeouts. Pattern:
#   1. start_* tool -> spawns subprocess, returns job_id immediately
#   2. status_* tool -> returns running/completed/failed + recent output
#   3. result_* tool -> full (possibly truncated) output if complete
#   4. cancel_job -> attempt to terminate running subprocess
# State is stored in an in-memory dict (not persistent across process restarts).
# ---------------------------------------------------------------------------

# DEFINE BUILD TOOLS

@mcp.tool()
async def ube_build(workspace_path: str = DEFAULT_WORKSPACE, verbose: bool = False) -> str:
    """Run a blocking UBE build and return full stdout/stderr.

    Parameters:
      workspace_path: Root directory for the build (-r parameter to UBE)
      verbose: If True adds -v for verbose UBE output
    """
    # Blocking build execution
    # Early existence check for UBE executable
    if not _ube_exists(UBE_PATH):
        return (
            f"UBE executable not found or not executable at '{UBE_PATH}' (source={UBE_PATH_SOURCE}). "
            "Set UBE_PATH environment variable to override."
        )
    try:
        args = ["UBE_BUILD", "-r", workspace_path] + (["-v"] if verbose else [])
        logger.info("Starting blocking UBE build subprocess: %s (verbose=%s)", workspace_path, verbose)
        process = await asyncio.create_subprocess_exec(
            UBE_PATH, *args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=workspace_path,
            stdin=asyncio.subprocess.DEVNULL,
        )
        logger.info("Blocking UBE build subprocess started (pid=%s)", process.pid)

        # Stream-read stdout and stderr concurrently to avoid any pipe buffer issues
        stdout_lines: list[str] = []
        stderr_lines: list[str] = []

        async def read_stream(stream, collector, label):
            try:
                while True:
                    line = await stream.readline()
                    if not line:
                        break
                    decoded = line.decode('utf-8', 'replace')
                    collector.append(decoded)
            except Exception as re:
                logger.exception("Error reading %s: %s", label, re)

        await asyncio.gather(
            read_stream(process.stdout, stdout_lines, 'stdout'),
            read_stream(process.stderr, stderr_lines, 'stderr')
        )
        returncode = await process.wait()
        logger.info("Blocking UBE build subprocess completed (pid=%s, returncode=%s)", process.pid, returncode)

        full_stdout = ''.join(stdout_lines)
        full_stderr = ''.join(stderr_lines)
        output = f"Exit code: {returncode}\n"
        output += f"STDOUT:\n{full_stdout}\n"
        if full_stderr:
            output += f"STDERR:\n{full_stderr}\n"
        return output
    except Exception as e:
        logger.exception("Blocking UBE build subprocess failed: %s", e)
        return f"Error running UBE build: {str(e)}"

@mcp.tool()
async def ube_clean(workspace_path: str = DEFAULT_WORKSPACE) -> str:
    """Run a blocking UBE clean and return full stdout/stderr."""
    if not _ube_exists(UBE_PATH):
        return (
            f"UBE executable not found or not executable at '{UBE_PATH}' (source={UBE_PATH_SOURCE}). "
            "Set UBE_PATH environment variable to override."
        )
    try:
        args = ["UBE_CLEAN", "-r", workspace_path, "-rm"]
        logger.info("(blocking) Starting UBE clean subprocess: %s", workspace_path)
        process = await asyncio.create_subprocess_exec(
            UBE_PATH, *args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=workspace_path,
            stdin=asyncio.subprocess.DEVNULL,
        )
        logger.info("(blocking) UBE clean subprocess started (pid=%s)", process.pid)
        stdout, stderr = await process.communicate()
        logger.info("(blocking) UBE clean subprocess completed (pid=%s, returncode=%s)", process.pid, process.returncode)
        output = f"Exit code: {process.returncode}\n"
        output += f"STDOUT:\n{stdout.decode('utf-8')}\n"
        if stderr:
            output += f"STDERR:\n{stderr.decode('utf-8')}\n"
        return output
    except Exception as e:
        logger.exception("(blocking) UBE clean subprocess failed: %s", e)
        return f"Error running UBE clean: {str(e)}"

@mcp.tool()
async def ube_compile_commands(workspace_path: str = DEFAULT_WORKSPACE) -> str:
    """Run a blocking UBE compile commands generation and return full stdout/stderr."""
    if not _ube_exists(UBE_PATH):
        return (
            f"UBE executable not found or not executable at '{UBE_PATH}' (source={UBE_PATH_SOURCE}). "
            "Set UBE_PATH environment variable to override."
        )
    try:
        args = ["UBE_COMPCMD", "-r", workspace_path]
        logger.info("(blocking) Starting UBE compile commands subprocess: %s", workspace_path)
        process = await asyncio.create_subprocess_exec(
            UBE_PATH, *args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=workspace_path,
            stdin=asyncio.subprocess.DEVNULL,
        )
        logger.info("(blocking) UBE compile commands subprocess started (pid=%s)", process.pid)
        stdout, stderr = await process.communicate()
        logger.info("(blocking) UBE compile commands subprocess completed (pid=%s, returncode=%s)", process.pid, process.returncode)
        output = f"Exit code: {process.returncode}\n"
        output += f"STDOUT:\n{stdout.decode('utf-8')}\n"
        if stderr:
            output += f"STDERR:\n{stderr.decode('utf-8')}\n"
        return output
    except Exception as e:
        logger.exception("(blocking) UBE compile commands subprocess failed: %s", e)
        return f"Error generating compile commands: {str(e)}"

@mcp.tool()
def check_build_status(workspace_path: str = DEFAULT_WORKSPACE) -> str:
    """Check if the project has been built successfully"""
    try:
        # Look for common build artifacts
        workspace = Path(workspace_path)
        build_artifacts = []
        
        # Common build output patterns
        patterns = ["*.elf", "*.hex"]
        
        for pattern in patterns:
            artifacts = list(workspace.rglob(pattern))
            if artifacts:
                build_artifacts.extend([str(a.relative_to(workspace)) for a in artifacts])
        
        if build_artifacts:
            return f"Build artifacts found:\n" + "\n".join(build_artifacts[:20])  # Limit output
        else:
            return "No build artifacts found. Project may not be built."
            
    except Exception as e:
        return f"Error checking build status: {str(e)}"

@mcp.tool()
def list_source_files(workspace_path: str = DEFAULT_WORKSPACE, extension: str = "c") -> str:
    """List source files in the workspace"""
    try:
        workspace = Path(workspace_path)
        source_files = list(workspace.rglob(f"*.{extension}"))
        
        if source_files:
            relative_paths = [str(f.relative_to(workspace)) for f in source_files]
            return f"Found {len(relative_paths)} {extension} files:\n" + "\n".join(relative_paths[:50])
        else:
            return f"No {extension} files found in workspace"
            
    except Exception as e:
        return f"Error listing source files: {str(e)}"

@mcp.tool()
def get_workspace_info(workspace_path: str = DEFAULT_WORKSPACE) -> str:
    """Get information about the workspace structure"""
    try:
        workspace = Path(workspace_path)
        if not workspace.exists():
            return f"Workspace path does not exist: {workspace_path}"
        
        info = f"Workspace: {workspace_path}\n"
        info += f"Exists: {workspace.exists()}\n"
        info += f"Is directory: {workspace.is_dir()}\n"
        
        if workspace.is_dir():
            # Count files by extension
            extensions = {}
            for file in workspace.rglob("*"):
                if file.is_file():
                    ext = file.suffix.lower()
                    extensions[ext] = extensions.get(ext, 0) + 1
            
            info += "\nFile types:\n"
            for ext, count in sorted(extensions.items()):
                if ext:  # Skip files without extension
                    info += f"  {ext}: {count} files\n"
        
        return info
    except Exception as e:
        return f"Error getting workspace info: {str(e)}"

# DEFINE RESOURCES

@mcp.resource("build-config://workspace")
def get_build_config() -> str:
    """Get the current build configuration"""
    config = {
        "ube_path": UBE_PATH,
        "ube_path_source": UBE_PATH_SOURCE,
        "platform": sys.platform,
        "os_name": os.name,
        "python_implementation": platform.python_implementation(),
        "default_workspace": DEFAULT_WORKSPACE,
        "available_commands": ["UBE_BUILD", "UBE_CLEAN", "UBE_COMPCMD"],
    }
    return json.dumps(config, indent=2)

@mcp.resource("build-log://{workspace}")
def get_build_log(workspace: str) -> str:
    """Get build log information for a workspace"""
    try:
        workspace_path = Path(workspace if workspace != "default" else DEFAULT_WORKSPACE)
        
        # Look for common log files
        log_files = []
        log_patterns = ["*.log", "build.txt", "compile.log", "error.log"]
        
        for pattern in log_patterns:
            logs = list(workspace_path.rglob(pattern))
            log_files.extend(logs)
        
        if log_files:
            result = "Build log files found:\n"
            for log_file in log_files[:10]:  # Limit to first 10 log files
                try:
                    with open(log_file, 'r', encoding='utf-8', errors='ignore') as f:
                        content = f.read()
                        result += f"\n=== {log_file.name} ===\n"
                        result += content[-1000:]  # Last 1000 characters
                        result += "\n"
                except Exception as e:
                    result += f"\nError reading {log_file}: {str(e)}\n"
            
            return result
        else:
            return f"No build log files found in {workspace_path}"
            
    except Exception as e:
        return f"Error accessing build logs: {str(e)}"

@mcp.resource("project-structure://{path}")
def get_project_structure(path: str = DEFAULT_WORKSPACE) -> str:
    """Get the project directory structure"""
    try:
        workspace = Path(path)
        structure = f"Project Structure for: {path}\n"
        structure += "=" * 50 + "\n"
        
        def build_tree(directory, prefix="", max_depth=3, current_depth=0):
            if current_depth >= max_depth:
                return ""
            
            items = []
            try:
                items = sorted(directory.iterdir())
            except PermissionError:
                return f"{prefix}[Permission Denied]\n"
            
            tree_str = ""
            for i, item in enumerate(items):
                is_last = i == len(items) - 1
                current_prefix = "└── " if is_last else "├── "
                tree_str += f"{prefix}{current_prefix}{item.name}\n"
                
                if item.is_dir() and current_depth < max_depth - 1:
                    next_prefix = prefix + ("    " if is_last else "│   ")
                    tree_str += build_tree(item, next_prefix, max_depth, current_depth + 1)
            
            return tree_str
        
        structure += build_tree(workspace)
        return structure
        
    except Exception as e:
        return f"Error getting project structure: {str(e)}"

def main() -> None:
    """Entry point for console scripts."""
    # Run the FastMCP server using stdio transport (newline-delimited JSON-RPC).
    # Any exceptions at top-level are logged to stderr; stdout remains protocol-clean.
    try:
        asyncio.run(mcp.run())
    except Exception as e:
        logger.exception("Fatal server error: %s", e)
        # Intentionally do not print to stdout; process will exit with non-zero code.

# execute and return the stdio output
if __name__ == "__main__":
    main()