param (
    [Parameter(Mandatory = $true)]
    [string]$ResourceGroup,
    
    [Parameter(Mandatory = $true)]
    [string]$SearchServiceName,
    
    [Parameter(Mandatory = $true)]
    [string]$AiFoundryName,
    
    [Parameter(Mandatory = $true)]
    [string]$StorageAccountName
)

Write-Host "Starting setup script..."

# Get user principal ID
Write-Host "Getting user principal ID..."
$userPrincipalId = az ad signed-in-user show `
    --query id `
    -o tsv
if ($LASTEXITCODE -ne 0) {
    Write-Error "Failed to get user principal ID (Exit code: $LASTEXITCODE)"
    exit $LASTEXITCODE
}
Write-Host "User Principal ID: $userPrincipalId"

# 1. Enable system-assigned managed identities for search service
Write-Host "Enabling system-assigned managed identity for Search Service..."
$searchIdentity = az search service update `
    --name $SearchServiceName `
    --resource-group $ResourceGroup `
    --identity-type SystemAssigned `
    --query identity.principalId `
    -o tsv
if ($LASTEXITCODE -ne 0) {
    Write-Error "Failed to enable managed identity for Search Service (Exit code: $LASTEXITCODE)"
    exit $LASTEXITCODE
}

# WARNING: The identity should already be enable for azure ai foundry.
Write-Host "Getting system-assigned managed identity principal id for AI Foundry..."
$aiFoundryIdentity = az cognitiveservices account identity show `
    --name $AiFoundryName `
    --resource-group $ResourceGroup `
    --query principalId `
    -o tsv
if ($LASTEXITCODE -ne 0) {
    Write-Error "Failed to get managed identity for AI Foundry (Exit code: $LASTEXITCODE)"
    exit $LASTEXITCODE
}

# 2. Get resource IDs
$searchServiceId = az search service show `
    --name $SearchServiceName `
    --resource-group $ResourceGroup `
    --query id `
    -o tsv
if ($LASTEXITCODE -ne 0) {
    Write-Error "Failed to get Search Service ID (Exit code: $LASTEXITCODE)"
    exit $LASTEXITCODE
}

$aiFoundryId = az cognitiveservices account show `
    --name $AiFoundryName `
    --resource-group $ResourceGroup `
    --query id `
    -o tsv
if ($LASTEXITCODE -ne 0) {
    Write-Error "Failed to get AI Foundry ID (Exit code: $LASTEXITCODE)"
    exit $LASTEXITCODE
}

$storageAccountId = az storage account show `
    --name $StorageAccountName `
    --resource-group $ResourceGroup `
    --query id `
    -o tsv
if ($LASTEXITCODE -ne 0) {
    Write-Error "Failed to get Storage Account ID (Exit code: $LASTEXITCODE)"
    exit $LASTEXITCODE
}

# 3. Assign roles on Azure Cognitive Search
Write-Host "Assigning roles on Search Service..."

# Search Index Data Reader to AI Foundry managed identity
$result = az role assignment create `
    --assignee-object-id $aiFoundryIdentity `
    --assignee-principal-type ServicePrincipal `
    --role 'Search Index Data Reader' `
    --scope $searchServiceId
if ($LASTEXITCODE -ne 0) {
    Write-Error "Failed to assign Search Index Data Reader role (Exit code: $LASTEXITCODE)"
    exit $LASTEXITCODE
}

# Search Service Contributor to AI Foundry managed identity
$result = az role assignment create `
    --assignee-object-id $aiFoundryIdentity `
    --assignee-principal-type ServicePrincipal `
    --role 'Search Service Contributor' `
    --scope $searchServiceId
if ($LASTEXITCODE -ne 0) {
    Write-Error "Failed to assign Search Service Contributor role (Exit code: $LASTEXITCODE)"
    exit $LASTEXITCODE
}

# Assign roles to user principal ID on Search Service
$roles = @(
    "Search Service Contributor",
    "Search Index Data Contributor"
)

foreach ($role in $roles) {
    Write-Host "Assigning $role role to user on Search Service..."
    $result = az role assignment create `
        --assignee-object-id $userPrincipalId `
        --assignee-principal-type User `
        --role $role `
        --scope $searchServiceId
    if ($LASTEXITCODE -ne 0) {
        Write-Error "Failed to assign $role role on Search Service (Exit code: $LASTEXITCODE)"
        exit $LASTEXITCODE
    }
}

# 4. Assign roles on Azure AI Foundry 
Write-Host "Assigning roles on AI Foundry..."

# Cognitive Services OpenAI Contributor to Search service managed identity
$result = az role assignment create `
    --assignee-object-id $searchIdentity `
    --assignee-principal-type ServicePrincipal `
    --role 'Cognitive Services OpenAI Contributor' `
    --scope $aiFoundryId
if ($LASTEXITCODE -ne 0) {
    Write-Error "Failed to assign Cognitive Services OpenAI Contributor role (Exit code: $LASTEXITCODE)"
    exit $LASTEXITCODE
}

# Assign roles to user principal ID on AI Foundry
$roles = @(
    "Cognitive Services OpenAI Contributor",
    "Cognitive Services Contributor",
    "Contributor"
)

foreach ($role in $roles) {
    Write-Host "Assigning $role role to user on AI Foundry..."
    $result = az role assignment create `
        --assignee-object-id $userPrincipalId `
        --assignee-principal-type User `
        --role $role `
        --scope $aiFoundryId
    if ($LASTEXITCODE -ne 0) {
        Write-Error "Failed to assign $role role on AI Foundry (Exit code: $LASTEXITCODE)"
        exit $LASTEXITCODE
    }
}

# 5. Assign roles on Azure Blob storage
Write-Host "Assigning roles on Storage Account..."

# Storage Blob Data Contributor to the AI Foundry managed identity
$result = az role assignment create `
    --assignee-object-id $aiFoundryIdentity `
    --assignee-principal-type ServicePrincipal `
    --role 'Storage Blob Data Contributor' `
    --scope $storageAccountId
if ($LASTEXITCODE -ne 0) {
    Write-Error "Failed to assign Storage Blob Data Contributor role (Exit code: $LASTEXITCODE)"
    exit $LASTEXITCODE
}

# Storage Blob Data Reader to the Search service managed identity
$result = az role assignment create `
    --assignee-object-id $searchIdentity `
    --assignee-principal-type ServicePrincipal `
    --role 'Storage Blob Data Reader' `
    --scope $storageAccountId
if ($LASTEXITCODE -ne 0) {
    Write-Error "Failed to assign Storage Blob Data Reader role (Exit code: $LASTEXITCODE)"
    exit $LASTEXITCODE
}

# Assign roles to user principal ID on Storage Account
$roles = @(
    "Contributor",
    "Storage Blob Data Contributor",
    "Storage File Data Privileged Contributor"
)

foreach ($role in $roles) {
    Write-Host "Assigning $role role to user on Storage Account..."
    $result = az role assignment create `
        --assignee-object-id $userPrincipalId `
        --assignee-principal-type User `
        --role $role `
        --scope $storageAccountId
    if ($LASTEXITCODE -ne 0) {
        Write-Error "Failed to assign $role role on Storage Account (Exit code: $LASTEXITCODE)"
        exit $LASTEXITCODE
    }
}

$resourceGroupId = az group show `
    --name $ResourceGroup `
    --query id `
    -o tsv
if ($LASTEXITCODE -ne 0) {
    Write-Error "Failed to get Resource Group ID (Exit code: $LASTEXITCODE)"
    exit $LASTEXITCODE
}

# 6. Assign Contributor role on Resource Group for web app deployment
Write-Host "Assigning Contributor role to user on Resource Group..."
$result = az role assignment create `
    --assignee-object-id $userPrincipalId `
    --assignee-principal-type User `
    --role 'Contributor' `
    --scope $resourceGroupId
if ($LASTEXITCODE -ne 0) {
    Write-Error "Failed to assign Contributor role on Resource Group (Exit code: $LASTEXITCODE)"
    exit $LASTEXITCODE
}

Write-Host "Role assignments completed successfully."