Parameters:
 EnvironmentName:
    Type: String
    Default: EKSBlogSetup
    Description: Environment name for resource naming and tagging

 ArgoCDLoadBalancerURL:
    Type: String
    Description: Load Balancer URL for the ArgoCD Server (example - https://example.us-east-1.elb.amazonaws.com) 

 AWSSecretName:
    Type: String
    Default: argocdcreds
    Description: AWS Secrets Manager Secret Name that was created to store ArgoCD Credentials

Resources:
  # Lambda Function 1: K8sGPT Agent Lambda
  LambdaK8sGPTAgent:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Sub "${EnvironmentName}-LambdaK8sGPTAgent"
      Handler: lambda_function.lambda_handler
      Role: !GetAtt LambdaK8sGPTAgentRole.Arn
      Runtime: python3.12
      Code:
        S3Bucket: aws-blogs-artifacts-public
        S3Key: ML-18112/lambda-k8sgpt-agent.zip
      Description: 'Lambda function that interacts with the K8sGPT Agent on the EKS cluster'
      Timeout: 30
      MemorySize: 256

  LambdaK8sGPTAgentRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: K8sGPTLambdaPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - eks:ListClusters
                  - eks:DescribeCluster
                Resource: "*"
              - Effect: Allow
                Action:
                  - ec2:DescribeInstances
                Resource: "*"

  # Lambda Function 2: Restart and Rollback Lambda
  RestartRollBackApplicationArgoCD:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Sub "${EnvironmentName}-RestartRollBackApplicationArgoCD"
      Handler: lambda_function.lambda_handler
      Role: !GetAtt RestartRollBackApplicationArgoCDRole.Arn
      Runtime: python3.12
      Code:
        S3Bucket: aws-blogs-artifacts-public
        S3Key: ML-18112/RestartRollBackApplicationArgoCD.zip
      Description: 'Lambda function for restarting and rolling back ArgoCD applications'
      Timeout: 30
      MemorySize: 256
      Environment:
        Variables:
          Region: !Ref "AWS::Region"
          ArgoCD_LoadBalancer_URL: !Ref ArgoCDLoadBalancerURL
          AWS_SECRET_NAME: !Ref AWSSecretName

  RestartRollBackApplicationArgoCDRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
        - arn:aws:iam::aws:policy/SecretsManagerReadWrite

  # Lambda Function 3: Memory Management Lambda
  ArgocdAdjustMemory:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Sub "${EnvironmentName}-ArgocdAdjustMemory"
      Handler: lambda_function.lambda_handler
      Role: !GetAtt ArgocdAdjustMemoryRole.Arn
      Runtime: python3.12
      Code:
        S3Bucket: aws-blogs-artifacts-public
        S3Key: ML-18112/argocdAdjustMemory.zip
      Description: 'Lambda function for managing ArgoCD application memory'
      Timeout: 30
      MemorySize: 256
      Environment:
        Variables:
          Region: !Ref "AWS::Region"
          ArgoCD_LoadBalancer_URL: !Ref ArgoCDLoadBalancerURL
          AWS_SECRET_NAME: !Ref AWSSecretName

  ArgocdAdjustMemoryRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
        - arn:aws:iam::aws:policy/SecretsManagerReadWrite

  # Bedrock Agent Execution Role
  BedrockAgentExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: bedrock.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: BedrockAgentCustomPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Sid: AmazonBedrockAgentBedrockFoundationModelPolicyProd
                Effect: Allow
                Action: bedrock:InvokeModel
                Resource:
                  - !Sub 'arn:aws:bedrock:${AWS::Region}::foundation-model/anthropic.claude-3-5-sonnet-20240620-v1:0'

  # K8sGPT Agent and Alias
  K8sGPTAgent:
    Type: AWS::Bedrock::Agent
    Properties:
      AgentName: !Sub "${EnvironmentName}-K8sGPT-Agent"
      Description: K8sGPT Analysis Agent for Kubernetes Cluster Diagnostics
      AutoPrepare: true
      Instruction: |
        You are a K8sGPT Analysis Agent that:

          1. Analysis Capabilities:
             - Retrieves cluster health status
             - Identifies potential issues and vulnerabilities
             - Provides detailed diagnostics of Kubernetes resources
             - Analyzes configuration and security concerns

          2. Response Protocol:
             - Always provides severity levels (Critical/Warning/Info)
             - Groups issues by resource types (Pods/Deployments/Nodes)
             - Includes timestamp of analysis
             - Lists specific error messages and their contexts

          3. Issue Reporting Format:
             Analysis Summary:
             - Cluster: [Cluster Name]
             - Time: [Analysis Timestamp]
             - Total Issues: [Count]

             Detailed Findings:
             - Resource: [Resource Type/Name]
             - Severityity: [Critical/Warning/Info]
             - Issue: [Problem Description]
             - Impact: [Potential Impact]
             - Recommendation: [Suggested Fix]

          4. Error Handling:
             - Reports connectivity issues clearly
             - Indicates when K8sGPT operator is unavailable
             - Provides troubleshooting steps for common errors

          5. Response Size Management:
             If the response exceeds 25KB:
             - Prioritize critical and warning issues over informational ones
             - Maintain the structured format while condensing detailed descriptions
             - Focus on high-impact issues first
             - Indicate that the response has been summarized
             - Provide guidance on how to get more detailed information if needed

          Note: Focuses on providing actionable insights and clear, structured analysis results.
      AgentResourceRoleArn: !GetAtt BedrockAgentExecutionRole.Arn
      IdleSessionTTLInSeconds: 1800
      FoundationModel: anthropic.claude-3-5-sonnet-20240620-v1:0
      ActionGroups:
        - ActionGroupName: k8s-cluster-operations
          Description: This action group can perform various operations on the EKS clusters, such as listing clusters, describing clusters, and executing other cluster-related actions.
          ActionGroupExecutor: 
            Lambda: !GetAtt LambdaK8sGPTAgent.Arn  # Changed from LambdaArn to Lambda
          FunctionSchema:
            Functions:
              - Name: "get_k8sgpt_results"
                Description: "Invokes Lambda for Kubernetes Cluster Diagnostics"
                RequireConfirmation: DISABLED
                Parameters:
                  eks_cluster:
                    Type: "string"
                    Description: "The name of the EKS cluster to analyze"
                    Required: true



  K8sGPTAgentAlias:
    Type: AWS::Bedrock::AgentAlias
    Properties:
      AgentAliasName: !Sub "${EnvironmentName}-K8sGPT"
      AgentId: !Ref K8sGPTAgent

  # ArgoCD Agent and Alias
  ArgoCDAgent:
      Type: AWS::Bedrock::Agent
      Properties:
        AgentName: !Sub "${EnvironmentName}-ArgoCD-Agent"
        Description: ArgoCD Operations Executor
        Instruction: |
          You are an ArgoCD operations executor that:

          1. Handles application management tasks:
             - Application restarts
             - Version rollbacks
             - Resource limit modifications

          2. Safety Protocols:
             - Always requires explicit user confirmation before executing changes
             - Presents confirmation message: "Please confirm [specific_action] for [application_name]. Reply with 'Confirm' to proceed or 'Deny' to cancel."
             - Validates application existence before operations

          3. Resource Management:
             - Modifies pod memory limits within allowed ranges
             - Implements changes gradually to prevent disruption

          4. Response Format:
             Operation Result:
             - Action: [Type of operation]
             - Target: [Application name]
             - Status: [Success/Failure]
             - Details: [Operation specifics]
             - Next Steps: [If any]

          5. Error Handling:
             - Provides clear error messages when operations fail
             - Suggests potential solutions for common errors

          6. Response Size Management:
             If the response exceeds 25KB:
             - Maintain critical operation details and status information
             - Summarize detailed logs while preserving essential error messages
             - Keep all safety-related information intact
             - Include a note indicating the response has been summarized
             - Provide information on how to access full operation details

          Note: When working in multi-agent mode, maintains same confirmation and response protocols.
        AutoPrepare: true
        AgentResourceRoleArn: !GetAtt BedrockAgentExecutionRole.Arn
        IdleSessionTTLInSeconds: 1800
        FoundationModel: anthropic.claude-3-5-sonnet-20240620-v1:0
        ActionGroups:
          - ActionGroupName: argocd-restart
            Description: Restarts specified application
            ActionGroupExecutor:
              Lambda: !GetAtt RestartRollBackApplicationArgoCD.Arn
            FunctionSchema:
              Functions:
                - Name: "APPLICATION_RESTART"
                  Description: "Restarts specified ArgoCD application with confirmation. Validates application existence and health status before restart. Requires explicit confirmation due to critical nature."
                  RequireConfirmation: DISABLED
                  Parameters:
                    app_name:
                      Type: "string"
                      Description: "ArgoCD application identifier"
                      Required: true
                    action_name:
                      Type: "string"
                      Description: "This values should always be Restart"
                      Required: true

          - ActionGroupName: argocd-rollback
            Description: Rolls back application to previous version
            ActionGroupExecutor:
              Lambda: !GetAtt RestartRollBackApplicationArgoCD.Arn
            FunctionSchema:
              Functions:
                - Name: "APPLICATION_ROLLBACK"
                  Description: "Rolls back ArgoCD application to previous version. Validates application existence and version availability. Requires explicit confirmation due to critical nature."
                  RequireConfirmation: DISABLED
                  Parameters:
                    app_name:
                      Type: "string"
                      Description: "ArgoCD application identifier"
                      Required: true
                    action_name:
                      Type: "string"
                      Description: "This values should always be Rollback"
                      Required: true

          - ActionGroupName: argocd-memory-management
            Description: Modifies pod memory limits for applications.
            ActionGroupExecutor:
              Lambda: !GetAtt ArgocdAdjustMemory.Arn
            FunctionSchema:
              Functions:
                - Name: "adjust_pod_memory"
                  Description: "adjusts pod memory limits for ArgoCD applications. Validates input parameters before execution. Memory limit must be in format [100Mi, 1Gi] where Mi=MB, Gi=GB. Performs gradual scaling with safety checks. Requires explicit confirmation due to critical nature."
                  RequireConfirmation: DISABLED
                  Parameters:
                    app_name:
                      Type: "string"
                      Description: "ArgoCD application identifier"
                      Required: true
                    memory_limit:
                      Type: "string"
                      Description: "Target memory limit. range: 100Mi-4Gi. examples: [\"500Mi\", \"1Gi\", \"2Gi\"]"
                      Required: true
                    pod_name:
                      Type: "string"
                      Description: "Target pod identifier in Argo CD application."
                      Required: true

  ArgoCDAgentAlias:
    Type: AWS::Bedrock::AgentAlias
    Properties:
      AgentAliasName: !Sub "${EnvironmentName}-ArgoCD"
      AgentId: !Ref ArgoCDAgent

  # IAM Policy for the collaborative Bedrock Agent
  CollabAgentPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Description: "Policy for Collaborative Bedrock Agent to invoke foundation models and collaborate with other agents"
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Sid: "AmazonBedrockAgentBedrockFoundationModelPolicyProd"
            Effect: "Allow"
            Action:
              - "bedrock:InvokeModel"
              - "bedrock:InvokeModelWithResponseStream"
            Resource:
              - !Sub "arn:aws:bedrock:${AWS::Region}::foundation-model/anthropic.claude-3-5-sonnet-20240620-v1:0"
          
          - Sid: "AmazonBedrockAgentsMultiAgentsPoliciesProd"
            Effect: "Allow"
            Action:
              - "bedrock:GetAgentAlias"
              - "bedrock:InvokeAgent"
            Resource:
              - !GetAtt K8sGPTAgentAlias.AgentAliasArn
              - !GetAtt ArgoCDAgentAlias.AgentAliasArn

  # IAM Role for the collaborator agent
  BedrockCollabAgentExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: bedrock.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - !Ref CollabAgentPolicy
      Description: "Role for Collaborative Bedrock Agent to invoke foundation models and collaborate with other agents"


  CollabBedrockAgent:
    Type: AWS::Bedrock::Agent
    Properties:
      AgentName: "EKSCollaboratorAgent"
      FoundationModel: "anthropic.claude-3-5-sonnet-20240620-v1:0"
      Instruction: |
        Supervisor Agent Instructions:

        You are a Supervisor Agent that coordinates between K8sGPT Analysis Agent and ArgoCD Agent to manage an EKS cluster. Follow these steps:

        1. Initial Analysis:
          a. Request cluster analysis from K8sGPT Agent:
              - Provide EKS_CLUSTER_NAME
              - Request "analyze_cluster" action
          b. Parse K8sGPT response, focusing on:
              - Severity levels (Critical/Warning/Info)
              - Resource types affected (Pods/Deployments/Nodes)
              - Specific error messages and contexts

        2. Issue Prioritization:
          a. Prioritize issues based on:
              - K8sGPT severity levels
              - Resource impact
          b. Focus on Critical and Warning level issues first

        3. Remediation Planning:
          For each prioritized issue:
          a. Determine appropriate ArgoCD action:
              - For application issues: restart_application or rollback_application
              - For resource issues: increase_pod_memory
          b. Prepare action parameters:
              - app_name
              - memory_limit (if applicable)
              - pod_name (if applicable)

        4. User Confirmation:
          Present structured recommendation:
          Issue: {detected_problem}
          Severity: {issue_severity}
          Proposed Fix: {argocd_action}
          Parameters: {action_parameters}
          Confirm action? [Confirm/Deny]

        5. Action Execution:
          If user confirms:
          a. Send to ArgoCD agent:
              - action: {selected_action}
              - parameters: {validated_params}
          b. Wait for ArgoCD agent response

        6. Result Verification:
          a. Check ArgoCD operation result
          b. Verify improvement in affected resources
          c. Prepare complete status report

        7. Result Reporting:
          Provide summary including:
          - Action executed
          - Target application/resource
          - Success/Failure status
          - Any error messages or next steps

        8. Error Handling:
          - If K8sGPT reports connectivity issues, retry analysis
          - If ArgoCD action fails, report error and suggested solutions

        9. Response Size Management:
            If any agent response exceeds 25KB:
            - Summarize the most critical information
            - Maintain structured format
            - Indicate that response has been summarized

        Always maintain clear communication between agents and user, ensuring safety protocols are followed and providing actionable insights throughout the process.
      AgentResourceRoleArn: !GetAtt BedrockCollabAgentExecutionRole.Arn
      IdleSessionTTLInSeconds: 1800
      AutoPrepare: true
      AgentCollaboration: "SUPERVISOR"
      AgentCollaborators:
        - AgentDescriptor:
            AliasArn: !GetAtt K8sGPTAgentAlias.AgentAliasArn
          CollaborationInstruction: |
            You are a K8sGPT Analysis Agent that:

            1. Analysis Capabilities:
              - Retrieves cluster health status
              - Identifies potential issues and vulnerabilities
              - Provides detailed diagnostics of Kubernetes resources
              - Analyzes configuration and security concerns

            2. Response Protocol:
              - Always provides severity levels (Critical/Warning/Info)
              - Groups issues by resource types (Pods/Deployments/Nodes)
              - Includes timestamp of analysis
              - Lists specific error messages and their contexts

            3. Issue Reporting Format:
              Analysis Summary:
              - Cluster: [Cluster Name]
              - Time: [Analysis Timestamp]
              - Total Issues: [Count]

              Detailed Findings:
              - Resource: [Resource Type/Name]
              - Severityity: [Critical/Warning/Info]
              - Issue: [Problem Description]
              - Impact: [Potential Impact]
              - Recommendation: [Suggested Fix]

            4. Error Handling:
              - Reports connectivity issues clearly
              - Indicates when K8sGPT operator is unavailable
              - Provides troubleshooting steps for common errors

            5. Response Size Management:
              If the response exceeds 25KB:
              - Prioritize critical and warning issues over informational ones
              - Maintain the structured format while condensing detailed descriptions
              - Focus on high-impact issues first
              - Indicate that the response has been summarized
              - Provide guidance on how to get more detailed information if needed

            Note: Focuses on providing actionable insights and clear, structured analysis results.
          CollaboratorName: "K8sGPTAgent"
          RelayConversationHistory: "TO_COLLABORATOR"
        
        - AgentDescriptor:
            AliasArn: !GetAtt ArgoCDAgentAlias.AgentAliasArn
          CollaborationInstruction: |
            You are an ArgoCD operations executor that:

            1. Handles application management tasks:
              - Application restarts
              - Version rollbacks
              - Resource limit modifications

            2. Safety Protocols:
              - Always requires explicit user confirmation before executing changes
              - Presents confirmation message: "Please confirm [specific_action] for [application_name]. Reply with 'Confirm' to proceed or 'Deny' to cancel."
              - Validates application existence before operations

            3. Resource Management:
              - Modifies pod memory limits within allowed ranges
              - Implements changes gradually to prevent disruption

            4. Response Format:
              Operation Result:
              - Action: [Type of operation]
              - Target: [Application name]
              - Status: [Success/Failure]
              - Details: [Operation specifics]
              - Next Steps: [If any]

            5. Error Handling:
              - Provides clear error messages when operations fail
              - Suggests potential solutions for common errors

            6. Response Size Management:
              If the response exceeds 25KB:
              - Maintain critical operation details and status information
              - Summarize detailed logs while preserving essential error messages
              - Keep all safety-related information intact
              - Include a note indicating the response has been summarized
              - Provide information on how to access full operation details

            Note: When working in multi-agent mode, maintains same confirmation and response protocols.
          CollaboratorName: "ArgoCDAgent"
          RelayConversationHistory: "TO_COLLABORATOR"

  CollabBedrockAgentAlias:
    Type: AWS::Bedrock::AgentAlias
    Properties:
      AgentAliasName: !Sub "${EnvironmentName}-ACollabBedrockAgent"
      AgentId: !Ref CollabBedrockAgent

  # Lambda Permissions
  LambdaK8sGPTAgentLambdaPermission:
    Type: AWS::Lambda::Permission
    DependsOn: K8sGPTAgent
    Properties:
      FunctionName: !Ref LambdaK8sGPTAgent
      Action: lambda:InvokeFunction
      Principal: bedrock.amazonaws.com
      SourceArn: !GetAtt K8sGPTAgent.AgentArn

  RestartRollBackApplicationArgoCDLambdaPermission:
    Type: AWS::Lambda::Permission
    DependsOn: ArgoCDAgent
    Properties:
      FunctionName: !Ref RestartRollBackApplicationArgoCD
      Action: lambda:InvokeFunction
      Principal: bedrock.amazonaws.com
      SourceArn: !GetAtt ArgoCDAgent.AgentArn

  ArgocdAdjusteMemoryLambdaPermission:
    Type: AWS::Lambda::Permission
    DependsOn: ArgoCDAgent
    Properties:
      FunctionName: !Ref ArgocdAdjustMemory
      Action: lambda:InvokeFunction
      Principal: bedrock.amazonaws.com
      SourceArn: !GetAtt ArgoCDAgent.AgentArn

Outputs:
  K8sGPTAgentAliasId:
    Description: ID of the K8sGPT Bedrock Agent Alias
    Value: !Ref K8sGPTAgentAlias
    Export:
      Name: !Sub "${EnvironmentName}-K8sGPTAgentAliasId"

  ArgoCDAgentAliasId:
    Description: ID of the ArgoCD Bedrock Agent Alias
    Value: !Ref ArgoCDAgentAlias
    Export:
      Name: !Sub "${EnvironmentName}-ArgoCDAgentAliasId"

  CollabBedrockAgentAliasId:
    Description: ID of the Collaborator Bedrock Agent Alias
    Value: !Ref CollabBedrockAgentAlias
    Export:
      Name: !Sub "${EnvironmentName}-CollabBedrockAgentAliasId"

  LambdaK8sGPTAgentRole:
    Description: Lambda Execution Role for the K8sGPT Lambda
    Value: !GetAtt LambdaK8sGPTAgentRole.Arn
