MCP Sandbox

from datetime import timedelta
from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordRequestForm

from mcp_sandbox.auth.auth import authenticate_user, get_current_active_user
from mcp_sandbox.auth.utils import (
    ACCESS_TOKEN_EXPIRE_MINUTES,
    create_access_token,
    get_password_hash,
    generate_api_key,
)
from mcp_sandbox.db.database import db
from mcp_sandbox.models.user import User, UserCreate, Token

router = APIRouter(prefix="/api", tags=["auth"])


@router.post("/register", response_model=User)
async def register_user(user_data: UserCreate):
    """Register a new user"""
    # Check if user already exists
    existing_user = db.get_user(username=user_data.username)
    if existing_user:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Username already registered",
        )

    existing_email = db.get_user(email=user_data.email)
    if existing_email:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail="Email already registered"
        )

    # Create new user with hashed password
    hashed_password = get_password_hash(user_data.password)
    api_key = generate_api_key()

    user_dict = {
        "username": user_data.username,
        "email": user_data.email,
        "hashed_password": hashed_password,
        "api_key": api_key,
    }

    # Create user in database
    new_user = db.create_user(user_dict)

    # Return user without hashed password
    return User(**new_user)


@router.post("/token", response_model=Token)
async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()):
    """Get access token for login"""
    user = authenticate_user(form_data.username, form_data.password)
    if not user:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect username or password",
            headers={"WWW-Authenticate": "Bearer"},
        )

    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={"sub": user.username}, expires_delta=access_token_expires
    )

    return {"access_token": access_token, "token_type": "bearer"}


@router.get("/users/me", response_model=User)
async def read_users_me(current_user: User = Depends(get_current_active_user)):
    """Get current logged in user"""
    return current_user


@router.get("/users/me/api-key")
async def get_api_key(current_user: User = Depends(get_current_active_user)):
    """Get API key for current user"""
    user = db.get_user(user_id=current_user.id)
    if not user or not user.get("api_key"):
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="API key not found"
        )

    return {"api_key": user["api_key"]}


@router.post("/users/me/api-key/regenerate")
async def regenerate_api_key(current_user: User = Depends(get_current_active_user)):
    """Regenerate API key for current user"""
    new_api_key = generate_api_key()
    updated_user = db.update_user(current_user.id, {"api_key": new_api_key})

    if not updated_user:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="User not found"
        )

    return {"api_key": new_api_key}


@router.get("/users/me/sandboxes")
async def get_user_sandboxes(current_user: User = Depends(get_current_active_user)):
    """Get all sandboxes for current user"""
    sandboxes = db.get_user_sandboxes(current_user.id)
    return {"sandboxes": sandboxes}


@router.delete("/users/me/sandboxes/{sandbox_id}")
async def delete_user_sandbox(
    sandbox_id: str, current_user: User = Depends(get_current_active_user)
):
    """Delete a sandbox by ID (both database record and Docker container)"""
    # First, check if the sandbox exists and belongs to the current user
    if not db.is_sandbox_owner(current_user.id, sandbox_id):
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Sandbox not found or does not belong to this user",
        )

    # Get the sandbox record to find the Docker container ID
    sandbox_record = db.get_sandbox(sandbox_id)
    if not sandbox_record:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Sandbox not found in database",
        )

    # Create a new instance of SandboxManager to handle Docker containers
    from mcp_sandbox.core.sandbox_modules.manager import SandboxManager
    from mcp_sandbox.utils.config import logger

    sandbox_manager = SandboxManager()

    # Get the Docker container ID from the record
    docker_container_id = sandbox_record.get("docker_container_id")
    if docker_container_id:
        logger.info(
            f"Deleting Docker container with ID: {docker_container_id} for sandbox: {sandbox_id}"
        )
        # Delete the Docker container using the Docker container ID
        result = sandbox_manager.delete_sandbox(docker_container_id)
        if not result.get("success", False):
            # If Docker deletion fails, log the error but continue to remove database record
            logger.error(
                f"Failed to delete Docker container: {result.get('message', 'Unknown error')}"
            )
    else:
        logger.warning(f"No Docker container ID found for sandbox: {sandbox_id}")

    # Delete the sandbox from the database
    if not db.delete_sandbox(sandbox_id):
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Failed to delete sandbox from database",
        )

    return {"message": "Sandbox deleted successfully"}