#!/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")

# Stores global configuration like user_id, repo_url, private_key_path, workspace_dir
config_store: Dict[str, str] = {} 

# Default API base URL
DEFAULT_API_BASE = "https://aicodinggym.com/api"
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", "repo_url", "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", "repo_url", "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 resolve_private_key_path(user_id: str) -> Optional[str]:
    """Resolve the user's private key path from config or default location."""
    config_path = config_store.get("private_key_path")
    if config_path and Path(config_path).exists():
        return config_path
    key_path = Path.home() / ".mcp-keys" / f"{user_id}_id_rsa"
    if key_path.exists():
        return str(key_path)
    return None


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


def add_to_gitignore(target_path: Path, patterns: list[str]) -> None:
    """Add ignore patterns to a gitignore-style file without duplicates."""
    existing: set[str] = set()
    if target_path.exists():
        existing = set(line.rstrip("\n") for line in target_path.read_text().splitlines())
    else:
        target_path.parent.mkdir(parents=True, exist_ok=True)

    to_add = [pattern for pattern in patterns if pattern not in existing]
    if not to_add:
        return

    content = ""
    if target_path.exists():
        content = target_path.read_text()
        if content and not content.endswith("\n"):
            content += "\n"
    content += "\n".join(to_add) + "\n"
    target_path.write_text(content)


def ignore_tests_dirs(problem_dir: Path) -> None:
    """
    Ignore all */tests directories locally and prevent tracking their changes.
    This updates .git/info/exclude, updates .gitignore if tracked, and marks
    tracked files under tests/ as skip-worktree.
    """
    ignore_patterns = ["**/tests/"]

    git_info_exclude = problem_dir / ".git" / "info" / "exclude"
    add_to_gitignore(git_info_exclude, ignore_patterns)

    gitignore_tracked = run_git_command("git ls-files .gitignore", str(problem_dir))
    if gitignore_tracked.returncode == 0 and gitignore_tracked.stdout.strip():
        gitignore_path = problem_dir / ".gitignore"
        add_to_gitignore(gitignore_path, ignore_patterns)
        run_git_command("git update-index --skip-worktree -- .gitignore", str(problem_dir))

    # Mark tracked files under any tests/ directory as skip-worktree.
    list_result = run_git_command("git ls-files -z", str(problem_dir))
    if list_result.returncode != 0:
        return
    for rel_path in filter(None, list_result.stdout.split("\0")):
        if "tests" in Path(rel_path).parts:
            quoted = shlex.quote(rel_path)
            run_git_command(f"git update-index --skip-worktree -- {quoted}", str(problem_dir))


# Load any persisted configuration on startup
load_config_store()


# Tool implementations
@mcp.tool()
def configure(
    user_id: str,
    workspace_dir: str
) -> 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
    :type workspace_dir: str
    :return: Configuration status message
    :rtype: str
    """
    api_base = DEFAULT_API_BASE
    try:
        if not workspace_dir or not workspace_dir.strip():
            return "Error: workspace_dir is required. Please provide a workspace path."

        # Generate SSH key pair
        private_key_path, public_key = generate_ssh_key_pair(user_id)
        
        # Register public key with server
        api_endpoint = f"{api_base}/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_url = data.get("repo_url")
        repo_name = data.get("repo_name")
        if not repo_url:
            return "Error: Server did not return a repository URL"
        
        # Resolve and store workspace directory
        resolved_workspace_dir = str(Path(workspace_dir).expanduser().resolve())

        # Store configuration
        config_store["user_id"] = user_id
        if repo_name:
            config_store["repo_name"] = repo_name
        config_store["repo_url"] = repo_url
        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 or repo_url}
Repository URL: {repo_url}
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.
    You must not perform any follow-up actions after fetching the problem, 
    such as installing dependencies or running tests. 
    Instead, you should wait for user's instructions or request explicit approval 
    from the user for the next steps.
    
    :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_url = config_store.get("repo_url")
        
        if not user_id:
            key_ids = list_user_key_ids()
            if len(key_ids) > 1:
                return (
                    "Error: Multiple keys found in ~/.mcp-keys. "
                    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_url:
            return "Error: Repository URL not configured. Please run /configure first."
        
        # Call backend /fetch-problem endpoint to ensure branch exists
        api_base = DEFAULT_API_BASE
        api_endpoint = f"{api_base}/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 info from response
        branch = data.get("branch_name", problem_id)
        fetch_message = data.get("message", "")

        # Get user's private key path
        private_key_path = resolve_private_key_path(user_id)
        if not private_key_path:
            return "Error: SSH key not found. Please run /configure first."

        # 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}"
                remove_github_dir(problem_dir)
                ignore_tests_dirs(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)
        ignore_tests_dirs(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,
    commit_message: str,
    workspace_dir: str = None,
    user_id: str = None,
    force_push: bool = False
) -> 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 commit_message: Custom commit message
    :type commit_message: str
    :param workspace_dir: Local directory where the repository was cloned (optional)
    :type workspace_dir: str
    :param user_id: Your user ID (optional if configured globally via /configure)
    :type user_id: str
    :param force_push: Force push the branch (use with caution)
    :type force_push: bool
    :return: Submission status message
    :rtype: str
    """
    try:
        commit_message = (commit_message or "").strip()
        # Use configured values if not provided
        user_id = user_id or config_store.get("user_id")
        
        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."
        
        workspace_dir = workspace_dir or config_store.get("workspace_dir", ".")
        # Find problem directory
        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 = resolve_private_key_path(user_id)
        if not private_key_path:
            return "Error: SSH key path not found. Please run /fetch again."
        
        key_path = Path(private_key_path)

        branch_result = run_git_command("git rev-parse --abbrev-ref HEAD", str(problem_dir))
        if branch_result.returncode != 0:
            return f"Error: Git branch detection failed:\n{branch_result.stderr}"
        branch = branch_result.stdout.strip()
        if not branch or branch == "HEAD":
            return f"Error: Detached HEAD. Please checkout branch '{problem_id}' before submitting."
        if branch != problem_id:
            return f"Error: Expected branch '{problem_id}', but current branch is '{branch}'. Please checkout the correct branch."

        # Refresh remote refs before deciding what to push
        fetch_result = run_git_command(f"git fetch origin {branch}", str(problem_dir), key_path)
        if fetch_result.returncode != 0:
            return f"Error: Git fetch failed:\n{fetch_result.stderr}"
        
        # 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():
            ahead_result = run_git_command(
                f"git rev-list --count origin/{branch}..HEAD",
                str(problem_dir)
            )
            if ahead_result.returncode != 0:
                return f"Error: Git rev-list failed:\n{ahead_result.stderr}"
            try:
                ahead_count = int(ahead_result.stdout.strip() or "0")
            except ValueError:
                ahead_count = 0
            if ahead_count <= 0:
                return "Warning: No changes to commit and no local commits to push. Your working directory is clean."

            # No new changes, but local commits are ahead of remote.
            hash_result = run_git_command("git rev-parse HEAD", str(problem_dir))
            if hash_result.returncode != 0:
                return f"Error: Git rev-parse failed:\n{hash_result.stderr}"
            commit_hash = hash_result.stdout.strip()

            message_result = run_git_command("git log -1 --format=%s", str(problem_dir))
            head_message = message_result.stdout.strip()
            if head_message:
                commit_message = head_message
        else:
            # Commit changes
            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
        if force_push:
            push_cmd = f"git push --force-with-lease origin {branch}"
        else:
            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_API_BASE}/submissions"
        payload = {
            "problem_id": problem_id,
            "user_id": user_id,
            "commit_hash": commit_hash,
            "branch": branch,
            "commit_message": commit_message,
            "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)}"


@mcp.tool()
def reset_problem(
    problem_id: str,
    user_id: str = None,
    workspace_dir: str = None
) -> str:
    """
    Restore a problem repository's content back to the original setup commit.
    This preserves history by creating a new commit when changes are applied,
    and removes untracked files.
    
    :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 workspace_dir: Local directory where the repository was cloned (optional)
    :type workspace_dir: str
    :return: Reset status message
    :rtype: str
    """
    try:
        user_id = user_id or config_store.get("user_id")
        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."

        workspace_dir = 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}"

        log_result = run_git_command("git log --format=%H:%s --reverse", str(problem_dir))
        if log_result.returncode != 0:
            return f"Error: Git log failed:\n{log_result.stderr}"

        setup_prefix = "Setup SWE-bench instance:"
        setup_commit = None
        for line in log_result.stdout.splitlines():
            parts = line.split(":", 1)
            if len(parts) != 2:
                continue
            commit_hash, subject = parts[0].strip(), parts[1].strip()
            if subject.startswith(setup_prefix):
                setup_commit = commit_hash
                break

        if not setup_commit:
            return (
                "Error: Could not find the original setup commit. "
                f"Expected a commit message starting with '{setup_prefix}'."
            )

        restore_cmd = f"git restore --source {setup_commit} --staged --worktree -- ."
        restore_result = run_git_command(restore_cmd, str(problem_dir))
        if restore_result.returncode != 0:
            return f"Error: Git restore failed:\n{restore_result.stderr}"

        clean_result = run_git_command("git clean -fd", str(problem_dir))
        if clean_result.returncode != 0:
            return f"Error: Git clean failed:\n{clean_result.stderr}"

        status_result = run_git_command("git status --porcelain", str(problem_dir))
        if status_result.returncode != 0:
            return f"Error: Git status failed:\n{status_result.stderr}"

        if status_result.stdout.strip():
            add_result = run_git_command("git add -A -- . ':(exclude).github'", str(problem_dir))
            if add_result.returncode != 0:
                return f"Error: Git add failed:\n{add_result.stderr}"

            commit_message = f"Restore content to setup state ({setup_commit[:8]})"
            commit_cmd = f'git commit -m "{commit_message}"'
            commit_result = run_git_command(commit_cmd, str(problem_dir))
            if commit_result.returncode != 0:
                return f"Error: Git commit failed:\n{commit_result.stderr}"

            new_hash_result = run_git_command("git rev-parse HEAD", str(problem_dir))
            if new_hash_result.returncode != 0:
                return f"Error: Git rev-parse failed:\n{new_hash_result.stderr}"
            new_commit = new_hash_result.stdout.strip()[:8]

            return f"""Repository content restored to original setup state for {problem_id}

Setup commit: {setup_commit[:8]}
New commit: {new_commit}
Content restored without rewriting history. Untracked files removed.
"""
        return f"""Repository content already matches the original setup state for {problem_id}

Setup commit: {setup_commit[:8]}
No changes to commit. Untracked files removed.
"""
    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()
