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

import os
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"


# Helper functions for Git SSH key management
def infer_user_id_from_keys() -> Optional[str]:
    """
    Infer user_id by scanning for SSH key files in ~/.mcp-keys/.
    Returns the user_id if found, None otherwise.
    """
    key_dir = Path.home() / ".mcp-keys"
    if not key_dir.exists():
        return None
    
    # Look for *_id_rsa files
    key_files = list(key_dir.glob("*_id_rsa"))
    if not key_files:
        return None
    
    # Extract user_id from the first key file found
    key_file = key_files[0]
    user_id = key_file.stem.replace("_id_rsa", "")
    return user_id


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


# 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"
        
        # Store configuration
        config_store["user_id"] = user_id
        config_store["repo_name"] = repo_name
        config_store["private_key_path"] = str(private_key_path)
        
        if workspace_dir:
            config_store["workspace_dir"] = workspace_dir
        
        return f"""Configuration saved successfully!

User ID: {user_id}
Repository: {repo_name}
Workspace: {workspace_dir or 'Default (./workspace)'}
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") or infer_user_id_from_keys()
        workspace_dir = workspace_dir or config_store.get("workspace_dir", "./workspace")
        repo_name = config_store.get("repo_name")
        
        if not user_id:
            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
        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}"
                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."
        
        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
) -> 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
    :return: Submission status message
    :rtype: str
    """
    try:
        # Use configured values if not provided
        user_id = user_id or config_store.get("user_id") or infer_user_id_from_keys()
        commit_message = commit_message or ""
        
        if not user_id:
            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 configured workspace or default)
        workspace_dir = config_store.get("workspace_dir", "./workspace")
        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
        result = run_git_command("git add -A", str(problem_dir))
        if result.returncode != 0:
            return f"Error: Git add failed:\n{result.stderr}"
        
        # Check if there are changes to commit
        status_result = run_git_command("git status --porcelain", 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()
