#!/usr/bin/env python3
"""
watsonx Orchestrate Lambda Deployer MCP Server

A FastMCP server that automates the creation and deployment of AWS Lambda 
functions with API Gateway REST APIs using Terraform.
"""

import os
import json
import subprocess
import shutil
from pathlib import Path
from typing import Optional, Dict, Any
from fastmcp import FastMCP

# Initialize FastMCP server
mcp = FastMCP("wxo-lambda-deployer")


def execute_command(command: str, cwd: Optional[str] = None) -> str:
    """Execute a shell command and return output."""
    try:
        result = subprocess.run(
            command,
            shell=True,
            cwd=cwd or os.getcwd(),
            capture_output=True,
            text=True,
            check=True
        )
        return result.stdout
    except subprocess.CalledProcessError as e:
        raise Exception(f"Command failed: {e.stderr or e.stdout}")


@mcp.tool()
def init_lambda_project(
    project_name: str,
    runtime: str = "nodejs20.x",
    handler: str = "handler",
    working_dir: Optional[str] = None
) -> str:
    """
    Initialize a new AWS Lambda project with API Gateway configuration.
    
    Args:
        project_name: Name of the Lambda project
        runtime: Lambda runtime (nodejs20.x, nodejs18.x, python3.12, python3.11)
        handler: Handler function name (default: handler)
        working_dir: Working directory for the project (optional)
    
    Returns:
        Success message with project details
    """
    base_dir = working_dir or os.getcwd()
    project_dir = os.path.join(base_dir, project_name)
    lambda_dir = os.path.join(project_dir, "lambda")
    terraform_dir = os.path.join(project_dir, "terraform")
    
    # Create directory structure
    os.makedirs(lambda_dir, exist_ok=True)
    os.makedirs(terraform_dir, exist_ok=True)
    
    # Determine file extension and create Lambda function
    is_node = runtime.startswith("nodejs")
    ext = "js" if is_node else "py"
    handler_file = "index" if is_node else "lambda_function"
    
    if is_node:
        lambda_code = f"""exports.{handler} = async (event) => {{
  console.log('Event:', JSON.stringify(event, null, 2));
  
  return {{
    statusCode: 200,
    headers: {{
      'Content-Type': 'application/json',
      'Access-Control-Allow-Origin': '*'
    }},
    body: JSON.stringify({{
      message: 'Hello from {project_name}!',
      timestamp: new Date().toISOString()
    }})
  }};
}};
"""
    else:
        lambda_code = f"""import json
from datetime import datetime

def {handler}(event, context):
    print('Event:', json.dumps(event))
    
    return {{
        'statusCode': 200,
        'headers': {{
            'Content-Type': 'application/json',
            'Access-Control-Allow-Origin': '*'
        }},
        'body': json.dumps({{
            'message': 'Hello from {project_name}!',
            'timestamp': datetime.now().isoformat()
        }})
    }}
"""
    
    with open(os.path.join(lambda_dir, f"{handler_file}.{ext}"), "w") as f:
        f.write(lambda_code)
    
    # Create Terraform configuration
    terraform_config = """terraform {
  required_version = ">= 1.0"
  
  required_providers {
    aws = {
      source  = "hashicorp/aws"
      version = "~> 5.0"
    }
    archive = {
      source  = "hashicorp/archive"
      version = "~> 2.0"
    }
    random = {
      source  = "hashicorp/random"
      version = "~> 3.0"
    }
  }
}

provider "aws" {
  region = var.aws_region
  
  default_tags {
    tags = {
      project = var.project_name
    }
  }
}
"""
    
    variables_config = f"""variable "aws_region" {{
  description = "AWS region for resources"
  type        = string
  default     = "us-east-1"
}}

variable "project_name" {{
  description = "Project name for resource tagging"
  type        = string
  default     = "{project_name}"
}}

variable "lambda_runtime" {{
  description = "Lambda runtime"
  type        = string
  default     = "{runtime}"
}}

variable "lambda_handler" {{
  description = "Lambda handler"
  type        = string
  default     = "{handler_file}.{handler}"
}}
"""
    
    main_config = """# S3 bucket for Lambda deployment package
resource "random_pet" "lambda_bucket_name" {
  prefix = var.project_name
  length = 2
}

resource "aws_s3_bucket" "lambda_bucket" {
  bucket        = random_pet.lambda_bucket_name.id
  force_destroy = true
}

resource "aws_s3_bucket_ownership_controls" "lambda_bucket" {
  bucket = aws_s3_bucket.lambda_bucket.id
  rule {
    object_ownership = "BucketOwnerPreferred"
  }
}

resource "aws_s3_bucket_acl" "lambda_bucket" {
  depends_on = [aws_s3_bucket_ownership_controls.lambda_bucket]
  bucket     = aws_s3_bucket.lambda_bucket.id
  acl        = "private"
}

# Archive Lambda function code
data "archive_file" "lambda_package" {
  type        = "zip"
  source_dir  = "${path.module}/../lambda"
  output_path = "${path.module}/lambda.zip"
}

# Upload Lambda package to S3
resource "aws_s3_object" "lambda_package" {
  bucket = aws_s3_bucket.lambda_bucket.id
  key    = "lambda.zip"
  source = data.archive_file.lambda_package.output_path
  etag   = data.archive_file.lambda_package.output_md5
}

# IAM role for Lambda
resource "aws_iam_role" "lambda_exec" {
  name = "${var.project_name}-lambda-role"

  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Action = "sts:AssumeRole"
      Effect = "Allow"
      Principal = {
        Service = "lambda.amazonaws.com"
      }
    }]
  })
}

resource "aws_iam_role_policy_attachment" "lambda_policy" {
  role       = aws_iam_role.lambda_exec.name
  policy_arn = "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
}

# CloudWatch Log Group
resource "aws_cloudwatch_log_group" "lambda_logs" {
  name              = "/aws/lambda/${var.project_name}"
  retention_in_days = 7
}

# Lambda Function
resource "aws_lambda_function" "main" {
  function_name    = var.project_name
  s3_bucket        = aws_s3_bucket.lambda_bucket.id
  s3_key           = aws_s3_object.lambda_package.key
  runtime          = var.lambda_runtime
  handler          = var.lambda_handler
  source_code_hash = data.archive_file.lambda_package.output_base64sha256
  role             = aws_iam_role.lambda_exec.arn
  timeout          = 30
  memory_size      = 256

  depends_on = [
    aws_cloudwatch_log_group.lambda_logs,
    aws_iam_role_policy_attachment.lambda_policy
  ]
}

# API Gateway
resource "aws_apigatewayv2_api" "main" {
  name          = "${var.project_name}-api"
  protocol_type = "HTTP"
  
  cors_configuration {
    allow_origins = ["*"]
    allow_methods = ["GET", "POST", "PUT", "DELETE", "OPTIONS"]
    allow_headers = ["*"]
  }
}

resource "aws_apigatewayv2_stage" "main" {
  api_id      = aws_apigatewayv2_api.main.id
  name        = "$default"
  auto_deploy = true

  access_log_settings {
    destination_arn = aws_cloudwatch_log_group.api_logs.arn
    format = jsonencode({
      requestId      = "$context.requestId"
      ip             = "$context.identity.sourceIp"
      requestTime    = "$context.requestTime"
      httpMethod     = "$context.httpMethod"
      routeKey       = "$context.routeKey"
      status         = "$context.status"
      protocol       = "$context.protocol"
      responseLength = "$context.responseLength"
    })
  }
}

resource "aws_cloudwatch_log_group" "api_logs" {
  name              = "/aws/apigateway/${var.project_name}"
  retention_in_days = 7
}

resource "aws_apigatewayv2_integration" "lambda" {
  api_id             = aws_apigatewayv2_api.main.id
  integration_type   = "AWS_PROXY"
  integration_uri    = aws_lambda_function.main.invoke_arn
  integration_method = "POST"
}

resource "aws_apigatewayv2_route" "default" {
  api_id    = aws_apigatewayv2_api.main.id
  route_key = "$default"
  target    = "integrations/${aws_apigatewayv2_integration.lambda.id}"
}

resource "aws_lambda_permission" "api_gateway" {
  statement_id  = "AllowAPIGatewayInvoke"
  action        = "lambda:InvokeFunction"
  function_name = aws_lambda_function.main.function_name
  principal     = "apigateway.amazonaws.com"
  source_arn    = "${aws_apigatewayv2_api.main.execution_arn}/*/*"
}
"""
    
    outputs_config = """output "api_endpoint" {
  description = "API Gateway endpoint URL"
  value       = aws_apigatewayv2_stage.main.invoke_url
}

output "lambda_function_name" {
  description = "Lambda function name"
  value       = aws_lambda_function.main.function_name
}

output "lambda_function_arn" {
  description = "Lambda function ARN"
  value       = aws_lambda_function.main.arn
}

output "s3_bucket_name" {
  description = "S3 bucket name for Lambda packages"
  value       = aws_s3_bucket.lambda_bucket.id
}
"""
    
    # Write Terraform files
    with open(os.path.join(terraform_dir, "terraform.tf"), "w") as f:
        f.write(terraform_config)
    with open(os.path.join(terraform_dir, "variables.tf"), "w") as f:
        f.write(variables_config)
    with open(os.path.join(terraform_dir, "main.tf"), "w") as f:
        f.write(main_config)
    with open(os.path.join(terraform_dir, "outputs.tf"), "w") as f:
        f.write(outputs_config)
    
    # Create README
    readme = f"""# {project_name}

AWS Lambda function with API Gateway REST API

## Structure

```
{project_name}/
├── lambda/           # Lambda function code
│   └── {handler_file}.{ext}
└── terraform/        # Infrastructure as Code
    ├── terraform.tf  # Provider configuration
    ├── variables.tf  # Input variables
    ├── main.tf       # Main resources
    └── outputs.tf    # Output values
```

## Deployment

1. Initialize Terraform:
   ```bash
   cd terraform
   terraform init
   ```

2. Review the plan:
   ```bash
   terraform plan
   ```

3. Deploy:
   ```bash
   terraform apply
   ```

4. Get the API endpoint:
   ```bash
   terraform output api_endpoint
   ```

## Testing

```bash
curl $(terraform output -raw api_endpoint)
```

## Cleanup

```bash
terraform destroy
```
"""
    
    with open(os.path.join(project_dir, "README.md"), "w") as f:
        f.write(readme)
    
    return f"""✅ Lambda project initialized successfully!

Project: {project_name}
Location: {project_dir}
Runtime: {runtime}
Handler: {handler_file}.{handler}

Structure created:
- lambda/{handler_file}.{ext} - Lambda function code
- terraform/ - Infrastructure configuration
- README.md - Documentation

Next steps:
1. cd {project_dir}/terraform
2. terraform init
3. terraform plan
4. terraform apply"""


@mcp.tool()
def deploy_lambda(
    project_dir: str,
    auto_approve: bool = False,
    aws_region: Optional[str] = None
) -> str:
    """
    Deploy Lambda function using Terraform (init, plan, apply).
    
    Args:
        project_dir: Path to the Lambda project directory
        auto_approve: Auto-approve Terraform apply (default: False)
        aws_region: AWS region override (optional)
    
    Returns:
        Deployment results and outputs
    """
    terraform_dir = os.path.join(project_dir, "terraform")
    
    if not os.path.exists(terraform_dir):
        raise Exception(f"Terraform directory not found: {terraform_dir}")
    
    results = []
    
    # Terraform init
    results.append("🔧 Initializing Terraform...")
    init_output = execute_command("terraform init", terraform_dir)
    results.append(init_output)
    
    # Terraform validate
    results.append("\n✓ Validating configuration...")
    validate_output = execute_command("terraform validate", terraform_dir)
    results.append(validate_output)
    
    # Terraform plan
    results.append("\n📋 Planning deployment...")
    plan_cmd = f"terraform plan -var='aws_region={aws_region}'" if aws_region else "terraform plan"
    plan_output = execute_command(plan_cmd, terraform_dir)
    results.append(plan_output)
    
    # Terraform apply
    if auto_approve:
        results.append("\n🚀 Deploying infrastructure...")
        apply_cmd = "terraform apply -auto-approve"
        if aws_region:
            apply_cmd += f" -var='aws_region={aws_region}'"
        apply_output = execute_command(apply_cmd, terraform_dir)
        results.append(apply_output)
        
        # Get outputs
        results.append("\n📊 Deployment outputs:")
        output_cmd = "terraform output -json"
        outputs = execute_command(output_cmd, terraform_dir)
        results.append(outputs)
    else:
        results.append("\n⚠️  Auto-approve not enabled. Run 'terraform apply' manually to deploy.")
    
    return "\n".join(results)


@mcp.tool()
def update_lambda_code(
    project_dir: str,
    code: str,
    auto_apply: bool = False
) -> str:
    """
    Update Lambda function code and optionally redeploy.
    
    Args:
        project_dir: Path to the Lambda project directory
        code: New Lambda function code
        auto_apply: Auto-apply changes (default: False)
    
    Returns:
        Update status message
    """
    lambda_dir = os.path.join(project_dir, "lambda")
    terraform_dir = os.path.join(project_dir, "terraform")
    
    if not os.path.exists(lambda_dir):
        raise Exception(f"Lambda directory not found: {lambda_dir}")
    
    # Find handler file
    files = os.listdir(lambda_dir)
    handler_file = next((f for f in files if f.endswith(('.js', '.py'))), None)
    
    if not handler_file:
        raise Exception("No Lambda handler file found")
    
    # Update the code
    with open(os.path.join(lambda_dir, handler_file), "w") as f:
        f.write(code)
    
    results = [f"✅ Updated {handler_file}"]
    
    if auto_apply:
        results.append("\n🔄 Redeploying Lambda function...")
        apply_output = execute_command("terraform apply -auto-approve", terraform_dir)
        results.append(apply_output)
    else:
        results.append("\n⚠️  Code updated. Run 'terraform apply' to deploy changes.")
    
    return "\n".join(results)


@mcp.tool()
def get_deployment_status(project_dir: str) -> str:
    """
    Get current deployment status and outputs.
    
    Args:
        project_dir: Path to the Lambda project directory
    
    Returns:
        Deployment status information
    """
    terraform_dir = os.path.join(project_dir, "terraform")
    
    if not os.path.exists(terraform_dir):
        raise Exception(f"Terraform directory not found: {terraform_dir}")
    
    # Check if Terraform is initialized
    tf_state_exists = os.path.exists(os.path.join(terraform_dir, ".terraform"))
    if not tf_state_exists:
        return "❌ Terraform not initialized. Run 'terraform init' first."
    
    results = ["📊 Deployment Status\n"]
    
    try:
        # Get Terraform state
        state_output = execute_command("terraform show -json", terraform_dir)
        state = json.loads(state_output)
        
        if state.get("values") and state["values"].get("outputs"):
            results.append("Outputs:")
            for key, value in state["values"]["outputs"].items():
                results.append(f"  {key}: {value.get('value')}")
        
        # Get resource count
        if state.get("values") and state["values"].get("root_module"):
            resources = state["values"]["root_module"].get("resources", [])
            results.append(f"\nResources: {len(resources)} deployed")
    except Exception:
        results.append("⚠️  No deployment found or state is empty")
    
    return "\n".join(results)


@mcp.tool()
def destroy_lambda(
    project_dir: str,
    auto_approve: bool = False
) -> str:
    """
    Destroy Lambda function and all associated resources.
    
    Args:
        project_dir: Path to the Lambda project directory
        auto_approve: Auto-approve destruction (default: False)
    
    Returns:
        Destruction status message
    """
    terraform_dir = os.path.join(project_dir, "terraform")
    
    if not os.path.exists(terraform_dir):
        raise Exception(f"Terraform directory not found: {terraform_dir}")
    
    results = ["🗑️  Destroying Lambda deployment..."]
    
    destroy_cmd = "terraform destroy"
    if auto_approve:
        destroy_cmd += " -auto-approve"
    
    try:
        destroy_output = execute_command(destroy_cmd, terraform_dir)
        results.append(destroy_output)
        results.append("\n✅ Deployment destroyed successfully")
    except Exception as e:
        results.append(f"\n❌ Destruction failed: {str(e)}")
    
    return "\n".join(results)


@mcp.tool()
def test_lambda_endpoint(
    project_dir: str,
    method: str = "GET",
    body: Optional[str] = None,
    path: str = "/"
) -> str:
    """
    Test the deployed Lambda API endpoint.
    
    Args:
        project_dir: Path to the Lambda project directory
        method: HTTP method (GET, POST, PUT, DELETE)
        body: Request body for POST/PUT requests (JSON string)
        path: API path (default: /)
    
    Returns:
        API response with status and timing
    """
    terraform_dir = os.path.join(project_dir, "terraform")
    
    # Get API endpoint from Terraform output
    output_cmd = "terraform output -json"
    outputs = json.loads(execute_command(output_cmd, terraform_dir))
    
    if "api_endpoint" not in outputs or not outputs["api_endpoint"].get("value"):
        raise Exception("API endpoint not found in Terraform outputs")
    
    endpoint = outputs["api_endpoint"]["value"] + path
    
    # Build curl command
    curl_cmd = f"curl -X {method} \"{endpoint}\""
    if body and method in ["POST", "PUT"]:
        curl_cmd += f" -H \"Content-Type: application/json\" -d '{body}'"
    curl_cmd += " -w '\\n\\nStatus: %{http_code}\\nTime: %{time_total}s'"
    
    results = [
        f"🧪 Testing endpoint: {endpoint}",
        f"Method: {method}\n"
    ]
    
    try:
        response = execute_command(curl_cmd)
        results.append("Response:")
        results.append(response)
    except Exception as e:
        results.append(f"❌ Request failed: {str(e)}")
    
    return "\n".join(results)


def main():
    """Main entry point for the MCP server."""
    mcp.run()


if __name__ == "__main__":
    main()

# Made with Bob
