#!/usr/bin/env python3
"""
MCP Server for AI Coding Gym
Provides /fetch and /submit tools for local problem solving workflow
"""

import json
import os
import shlex
import shutil
import subprocess
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Optional
import requests

from mcp.server.fastmcp import FastMCP


# Initialize FastMCP server
mcp = FastMCP("ai-coding-gym")

# Global state
credentials_store: Dict[str, Dict[str, Any]] = {}  # problem_id -> {deploy_key, repo_url, branch, etc}
config_store: Dict[str, str] = {}  # Global config: user_id, server_url, api_key

# Default server URL
DEFAULT_SERVER_URL = "http://ai-coding-gym.cis240515.projects.jetstream-cloud.org:4000"
CONFIG_PATH = Path.home() / ".mcp-keys" / "config.json"


# Helper functions for Git SSH key management
def list_user_key_ids() -> list[str]:
    """List user_ids inferred from ~/.mcp-keys/*_id_rsa."""
    key_dir = Path.home() / ".mcp-keys"
    if not key_dir.exists():
        return []
    key_files = list(key_dir.glob("*_id_rsa"))
    user_ids: list[str] = []
    for key_file in key_files:
        user_id = key_file.stem.replace("_id_rsa", "")
        if user_id:
            user_ids.append(user_id)
    return user_ids


def load_config_store() -> None:
    """Load persisted configuration from ~/.mcp-keys/config.json if present."""
    try:
        if not CONFIG_PATH.exists():
            return
        data = json.loads(CONFIG_PATH.read_text())
        if not isinstance(data, dict):
            return
        for key in ("user_id", "repo_name", "private_key_path", "workspace_dir"):
            val = data.get(key)
            if isinstance(val, str) and val:
                config_store[key] = val
    except Exception:
        # If config is invalid, ignore and proceed with in-memory defaults
        return


def save_config_store() -> None:
    """Persist configuration to ~/.mcp-keys/config.json."""
    key_dir = Path.home() / ".mcp-keys"
    key_dir.mkdir(mode=0o700, exist_ok=True)
    data = {
        key: value
        for key, value in config_store.items()
        if key in ("user_id", "repo_name", "private_key_path", "workspace_dir")
    }
    CONFIG_PATH.write_text(json.dumps(data, indent=2))




def generate_ssh_key_pair(user_id: str) -> tuple[Path, str]:
    """
    Generate an SSH key pair for the user.
    Returns tuple of (private_key_path, public_key_content)
    """
    key_dir = Path.home() / ".mcp-keys"
    key_dir.mkdir(mode=0o700, exist_ok=True)
    
    key_path = key_dir / f"{user_id}_id_rsa"
    
    # Generate SSH key pair if it doesn't exist
    if not key_path.exists():
        result = subprocess.run(
            ["ssh-keygen", "-t", "rsa", "-b", "4096", "-f", str(key_path), "-N", "", "-C", f"mcp-{user_id}"],
            capture_output=True,
            text=True
        )
        if result.returncode != 0:
            raise RuntimeError(f"Failed to generate SSH key: {result.stderr}")
    
    # Read public key
    pub_key_path = Path(str(key_path) + ".pub")
    public_key = pub_key_path.read_text().strip()
    
    return key_path, public_key


def setup_deploy_key(problem_id: str, deploy_key: str) -> Path:
    """
    Write deployment key to a temporary file with secure permissions.
    Returns the path to the key file.
    """
    key_dir = Path.home() / ".mcp-keys"
    key_dir.mkdir(mode=0o700, exist_ok=True)
    
    key_path = key_dir / f"{problem_id}_deploy_key"
    key_path.write_text(deploy_key)
    key_path.chmod(0o600)
    
    return key_path


def cleanup_deploy_key(problem_id: str) -> None:
    """Remove deployment key file for a problem."""
    key_path = Path.home() / ".mcp-keys" / f"{problem_id}_deploy_key"
    if key_path.exists():
        key_path.unlink()


def run_git_command(cmd: str, cwd: str, key_path: Optional[Path] = None) -> subprocess.CompletedProcess:
    """
    Execute a git command with optional SSH key configuration.
    """
    env = os.environ.copy()
    if key_path:
        env["GIT_SSH_COMMAND"] = f"ssh -i {key_path} -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null"
    
    result = subprocess.run(
        cmd,
        shell=True,
        cwd=cwd,
        capture_output=True,
        text=True,
        env=env
    )
    return result


def remove_github_dir(problem_dir: Path) -> Optional[str]:
    """
    Remove .github directory locally and mark its files as skip-worktree
    so deletions/modifications are not tracked.
    """
    github_dir = problem_dir / ".github"
    if not github_dir.exists():
        return None

    # Mark tracked .github files as skip-worktree to avoid tracking deletions
    list_result = run_git_command("git ls-files .github", str(problem_dir))
    if list_result.returncode == 0:
        for rel_path in filter(None, list_result.stdout.splitlines()):
            quoted = shlex.quote(rel_path)
            run_git_command(f"git update-index --skip-worktree -- {quoted}", str(problem_dir))

    shutil.rmtree(github_dir, ignore_errors=True)
    return None


# Load any persisted configuration on startup
load_config_store()


# Tool implementations
@mcp.tool()
def configure(
    user_id: str,
    workspace_dir: str = None
) -> str:
    """
    Configure global settings for the MCP server. Generates an SSH key pair locally, 
    sends the public key to the server, and receives your repository name. 
    Set your user_id once. These values will be used as defaults for all fetch and submit operations.
    
    :param user_id: Your user ID for authentication with the backend server
    :type user_id: str
    :param workspace_dir: Default workspace directory for cloning repositories (optional)
    :type workspace_dir: str
    :return: Configuration status message
    :rtype: str
    """
    server_url = DEFAULT_SERVER_URL
    try:
        # Generate SSH key pair
        private_key_path, public_key = generate_ssh_key_pair(user_id)
        
        # Register public key with server
        api_endpoint = f"{server_url}/api/configure"
        payload = {
            "user_id": user_id,
            "public_key": public_key
        }
        
        response = requests.post(api_endpoint, json=payload, timeout=30)
        response.raise_for_status()
        data = response.json()
        
        repo_name = data.get("repo_name")
        if not repo_name:
            return "Error: Server did not return a repository name"
        
        # Resolve and store workspace directory (default to server CWD)
        if workspace_dir:
            resolved_workspace_dir = str(Path(workspace_dir).resolve())
        else:
            resolved_workspace_dir = str(Path.cwd().resolve())

        # Store configuration
        config_store["user_id"] = user_id
        config_store["repo_name"] = repo_name
        config_store["private_key_path"] = str(private_key_path)
        config_store["workspace_dir"] = resolved_workspace_dir
        save_config_store()
        
        return f"""Configuration saved successfully!

User ID: {user_id}
Repository: {repo_name}
Workspace: {resolved_workspace_dir}
SSH Key: {private_key_path}

Your public key has been registered with the server.
You can now use /fetch and /submit without passing these parameters.
"""
    
    except requests.RequestException as e:
        return f"Error: Failed to register with server: {str(e)}"
    except Exception as e:
        return f"Error: Configuration failed: {str(e)}"


@mcp.tool()
def fetch_problem(
    problem_id: str,
    user_id: str = None,
    workspace_dir: str = None
) -> str:
    """
    Fetch a coding problem from the server and clone the repository locally. 
    Uses your SSH key from /configure to access the repository.
    
    :param problem_id: The unique identifier for the problem (e.g., 'django__django-10097')
    :type problem_id: str
    :param user_id: Your user ID (optional if configured globally via /configure)
    :type user_id: str
    :param workspace_dir: Local directory to clone the repository into (optional)
    :type workspace_dir: str
    :return: Fetch status message
    :rtype: str
    """
    try:
        # Use configured values if not provided
        user_id = user_id or config_store.get("user_id")
        workspace_dir = workspace_dir or config_store.get("workspace_dir", ".")
        repo_name = config_store.get("repo_name")
        
        if not user_id:
            key_ids = list_user_key_ids()
            if len(key_ids) > 1:
                return (
                    "Error: Multiple keys found in ~/.mcp-keys. "
                    "Please run /configure with one user_id to select it. "
                    f"Found: {', '.join(sorted(key_ids))}."
                )
            return "Error: Could not determine user_id. Please run /configure first to set up your credentials."
        
        if not repo_name:
            return "Error: Repository not configured. Please run /configure first."
        
        # Call backend /fetch-problem endpoint to ensure branch exists
        server_url = DEFAULT_SERVER_URL
        api_endpoint = f"{server_url}/api/fetch-problem"
        payload = {
            "user_id": user_id,
            "problem_id": problem_id
        }
        
        try:
            response = requests.post(api_endpoint, json=payload, timeout=30)
            response.raise_for_status()
            data = response.json()
        except requests.RequestException as e:
            return f"Error: Failed to fetch problem from server: {str(e)}"
        
        # Extract branch and repo info from response
        branch = data.get("branch_name", problem_id)
        repo_url = data.get("repo_url")
        fetch_message = data.get("message", "")
        
        if not repo_url:
            return f"Error: Server did not return repository URL"
        
        # Get user's private key path
        private_key_path = config_store.get("private_key_path")
        if not private_key_path:
            # Try to infer from user_id
            key_path = Path.home() / ".mcp-keys" / f"{user_id}_id_rsa"
            if not key_path.exists():
                return "Error: SSH key not found. Please run /configure first."
            private_key_path = str(key_path)
        
        # Store credentials for later use
        credentials_store[problem_id] = {
            "repo_url": repo_url,
            "branch": branch,
            "user_id": user_id,
            "private_key_path": private_key_path
        }
        
        # Use user's own SSH key
        key_path = Path(private_key_path)
        
        # Create workspace directory
        workspace_path = Path(workspace_dir).resolve()
        workspace_path.mkdir(parents=True, exist_ok=True)

        problem_dir = workspace_path / problem_id
        # Store resolved workspace path for submit convenience
        credentials_store[problem_id]["workspace_dir"] = str(workspace_path)
        if problem_dir.exists():
            # Check if it's already on the right branch
            check_branch = f"cd {problem_dir} && git rev-parse --abbrev-ref HEAD"
            result = subprocess.run(check_branch, shell=True, capture_output=True, text=True)
            if result.returncode == 0 and result.stdout.strip() == branch:
                # Pull latest changes
                pull_cmd = f"git pull origin {branch}"
                result = run_git_command(pull_cmd, str(problem_dir), key_path)
                if result.returncode != 0:
                    return f"Error: Git pull failed:\n{result.stderr}"
                remove_github_dir(problem_dir)
                return f"""Problem {problem_id} already exists. Updated to latest version.

Repository: {problem_dir}
Branch: {branch}
Server: {fetch_message}

You can continue working on your solution!
"""
            else:
                return f"Error: Directory {problem_dir} already exists with different content. Please remove it first or use a different workspace."
        
        # Clone only the specific branch (shallow clone for efficiency)
        clone_cmd = f"git clone --single-branch --branch {branch} --depth 1 {repo_url} {problem_id}"
        result = run_git_command(clone_cmd, str(workspace_path), key_path)
        
        if result.returncode != 0:
            return f"Error: Git clone failed:\n{result.stderr}\n\nMake sure the branch '{branch}' exists in the repository."

        remove_github_dir(problem_dir)
        
        return f"""Successfully fetched problem: {problem_id}

Repository cloned to: {problem_dir}
Branch: {branch}
Server: {fetch_message}

The repository contains only the problem branch (shallow clone to save space).
You can now start working on the solution!
"""
    
    except Exception as e:
        return f"Error: Unexpected error: {str(e)}"


@mcp.tool()
def submit_solution(
    problem_id: str,
    user_id: str = None,
    commit_message: str = None,
    workspace_dir: str = None
) -> str:
    """
    Submit your solution by committing all changes and pushing to the remote repository. 
    Notifies the backend server that a submission has been made.
    
    :param problem_id: The unique identifier for the problem
    :type problem_id: str
    :param user_id: Your user ID (optional if configured globally via /configure)
    :type user_id: str
    :param commit_message: Custom commit message (optional)
    :type commit_message: str
    :param workspace_dir: Local directory where the repository was cloned (optional)
    :type workspace_dir: str
    :return: Submission status message
    :rtype: str
    """
    try:
        # Use configured values if not provided
        user_id = user_id or config_store.get("user_id")
        commit_message = commit_message or ""
        
        if not user_id:
            key_ids = list_user_key_ids()
            if len(key_ids) > 1:
                return (
                    "Error: Multiple keys found in ~/.mcp-keys. "
                    "Please run /configure with one user_id to select it. "
                    f"Found: {', '.join(sorted(key_ids))}."
                )
            return "Error: Could not determine user_id. Please run /configure first to set up your credentials."
        
        # Check if we have cached credentials
        if problem_id not in credentials_store:
            return f"Error: No credentials found for {problem_id}. Please run /fetch first."
        
        creds = credentials_store[problem_id]
        
        # Verify user_id matches
        if creds["user_id"] != user_id:
            return f"Error: User ID mismatch. Problem was fetched by {creds['user_id']}, not {user_id}."
        
        # Find problem directory (use explicit, cached, or configured workspace)
        workspace_dir = (
            workspace_dir
            or creds.get("workspace_dir")
            or config_store.get("workspace_dir", ".")
        )
        workspace_path = Path(workspace_dir).resolve()
        problem_dir = workspace_path / problem_id
        
        if not problem_dir.exists():
            return f"Error: Problem directory not found at {problem_dir}"
        
        # Use user's SSH key
        private_key_path = creds.get("private_key_path") or config_store.get("private_key_path")
        if not private_key_path:
            return "Error: SSH key path not found. Please run /fetch again."
        
        key_path = Path(private_key_path)
        
        # Stage all changes except .github
        result = run_git_command("git add -A -- . ':(exclude).github'", str(problem_dir))
        if result.returncode != 0:
            return f"Error: Git add failed:\n{result.stderr}"
        
        # Check if there are staged changes to commit
        status_result = run_git_command("git diff --cached --name-only", str(problem_dir))
        if not status_result.stdout.strip():
            return "Warning: No changes to commit. Your working directory is clean."
        
        # Commit changes
        if not commit_message:
            commit_message = f"Solution submission for {problem_id} at {datetime.now().isoformat()}"
        
        commit_cmd = f'git commit -m "{commit_message}"'
        result = run_git_command(commit_cmd, str(problem_dir))
        if result.returncode != 0:
            return f"Error: Git commit failed:\n{result.stderr}"
        
        # Get commit hash
        hash_result = run_git_command("git rev-parse HEAD", str(problem_dir))
        commit_hash = hash_result.stdout.strip()
        
        # Push to remote
        branch = creds["branch"]
        push_cmd = f"git push origin {branch}"
        result = run_git_command(push_cmd, str(problem_dir), key_path)
        
        if result.returncode != 0:
            return f"Error: Git push failed:\n{result.stderr}"
        
        # Notify backend server
        api_endpoint = f"{DEFAULT_SERVER_URL}/api/submit"
        payload = {
            "problem_id": problem_id,
            "user_id": user_id,
            "commit_hash": commit_hash,
            "branch": branch,
            "timestamp": datetime.now().isoformat()
        }
        
        response = requests.post(api_endpoint, json=payload, timeout=30)
        response.raise_for_status()
        
        return f"""Successfully submitted solution for {problem_id}

Commit: {commit_hash[:8]}
Branch: {branch}
Pushed to remote
Backend notified

Your solution has been submitted for evaluation!
"""
    
    except requests.RequestException as e:
        return f"Warning: Changes committed and pushed, but failed to notify backend:\n{str(e)}"
    except Exception as e:
        return f"Error: Unexpected error: {str(e)}"


def main():
    """Entry point for the MCP server."""
    mcp.run(transport='stdio')


if __name__ == "__main__":
    main()
