import json
import pulumi
import pulumi_aws as aws

def create_roles_and_policies():
    # IAM Role (Job Scheduler) 
    role_job_scheduler = aws.iam.Role("assitente-produtos-servicos-role-role", 
        assume_role_policy=json.dumps({
            "Version": "2012-10-17",
            "Statement": [{
            "Sid": "AmazonBedrockKnowledgeBaseTrustPolicy",
            "Effect": "Allow",
            "Principal": {
                "Service": "bedrock.amazonaws.com"
            },
            "Action": "sts:AssumeRole",
            "Condition": {
                "StringEquals": {
                    "aws:SourceAccount": "277048801940"
                },
                "ArnLike": {
                    "aws:SourceArn": "arn:aws:bedrock:us-east-1:277048801940:knowledge-base/*"
                }
            }
        }]
        }),
    )

    aws.iam.RolePolicy("assistente-produtos-servicos-role-role-policy",
        role=role_job_scheduler.id,
        policy=json.dumps({
            "Version": "2012-10-17",
            "Statement": [
                {
            "Sid": "BedrockInvokeModelStatement",
            "Effect": "Allow",
            "Action": [
                "bedrock:InvokeModel"
            ],
            "Resource": [
                "arn:aws:bedrock:us-east-1::foundation-model/amazon.titan-embed-text-v2:0"
            ]
        },{
            "Sid": "S3ListBucketStatement",
            "Effect": "Allow",
            "Action": [
                "s3:ListBucket"
            ],
            "Resource": [
                "arn:aws:s3:::frente-corretora-assistente-produtos-servicos-doc-storage"
            ],
            "Condition": {
                "StringEquals": {
                    "aws:ResourceAccount": [
                        "277048801940"
                    ]
                }
            }
        },
        {
            "Sid": "S3GetObjectStatement",
            "Effect": "Allow",
            "Action": [
                "s3:GetObject"
            ],
            "Resource": [
                "arn:aws:s3:::frente-corretora-assistente-produtos-servicos-doc-storage/*"
            ],
            "Condition": {
                "StringEquals": {
                    "aws:ResourceAccount": [
                        "277048801940"
                    ]
                }
            }
        },{
            "Sid": "S3VectorsPermissions",
            "Effect": "Allow",
            "Action": [
                "s3vectors:GetIndex",
                "s3vectors:QueryVectors",
                "s3vectors:PutVectors",
                "s3vectors:GetVectors",
                "s3vectors:DeleteVectors"
            ],
            "Resource": "arn:aws:s3vectors:us-east-1:277048801940:bucket/bedrock-knowledge-base-icy5rp/index/bedrock-knowledge-base-default-index",
            "Condition": {
                "StringEquals": {
                    "aws:ResourceAccount": "277048801940"
                }
            }
        }
            ]
        }),
    )

    return role_job_scheduler