import json
import boto3
import requests
import re
ecs_client = boto3.client('ecs')
ecr_client = boto3.client('ecr')

url = 'https://hooks.slack.com/services/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'
headers = {
    'Content-Type': 'application/json'
}

def parse_codedeploy_appspec_yaml(yaml_content):
    """
    Simple parser for CodeDeploy AppSpec YAML content.
    Extracts TaskDefinition and Hooks information.
    """
    result = {
        'Resources': [],
        'Hooks': []
    }
    
    try:
        # Extract TaskDefinition ARN - fixed regex to properly capture quoted content
        task_def_pattern = r'TaskDefinition:\s*"([^"]*)"'
        task_def_match = re.search(task_def_pattern, yaml_content)
        if task_def_match:
            task_definition_arn = task_def_match.group(1).strip()
            print(f"Extracted TaskDefinition ARN: {task_definition_arn}")
            result['Resources'] = [{
                'TargetService': {
                    'Properties': {
                        'TaskDefinition': task_definition_arn
                    }
                }
            }]
        else:
            print("WARNING: Could not find TaskDefinition in YAML content")
            print(f"YAML content: {yaml_content}")
        
        # Extract Hooks - fixed regex patterns
        hooks_section = re.search(r'Hooks:\s*\n(.*?)(?:\n\S|\Z)', yaml_content, re.DOTALL)
        if hooks_section:
            hooks_content = hooks_section.group(1)
            print(f"Hooks content: {hooks_content}")
            # Find all hook entries with corrected regex
            hook_entries = [
                ('BeforeInstall', r'-\s*BeforeInstall:\s*"([^"]*)"'),
                ('BeforeAllowTraffic', r'-\s*BeforeAllowTraffic:\s*"([^"]*)"'),
                ('AfterAllowTraffic', r'-\s*AfterAllowTraffic:\s*"([^"]*)"'),
                ('ApplicationStop', r'-\s*ApplicationStop:\s*"([^"]*)"'),
                ('ApplicationStart', r'-\s*ApplicationStart:\s*"([^"]*)"')
            ]
            
            for hook_name, pattern in hook_entries:
                matches = re.findall(pattern, hooks_content)
                for match in matches:
                    result['Hooks'].append({hook_name: match.strip()})
        
        return result
    
    except Exception as e:
        print(f"Error parsing YAML: {e}")
        return result

def get_git_commit_from_ecr_image(image_uri):
    """
    Extract git commit hash from ECR image tags.
    Returns the git commit hash or the original tag if not found.
    """
    try:
        # Parse the image URI to get repository and tag/digest
        if '@sha256:' in image_uri:
            # Image referenced by digest
            repo_and_digest = image_uri.split('@sha256:')
            repository_name = repo_and_digest[0].split('/')[-1]
            image_digest = 'sha256:' + repo_and_digest[1]
            
            # Get image details by digest
            response = ecr_client.describe_images(
                repositoryName=repository_name,
                imageIds=[{'imageDigest': image_digest}]
            )
        else:
            # Image referenced by tag
            image_tag = image_uri.split(':')[-1]
            repository_name = image_uri.split(':')[0].split('/')[-1]
            
            # Get image details by tag
            response = ecr_client.describe_images(
                repositoryName=repository_name,
                imageIds=[{'imageTag': image_tag}]
            )
        
        if response['imageDetails']:
            all_tags = response['imageDetails'][0].get('imageTags', [])
            # Look for git commit hash tags (typically 7-40 characters, alphanumeric)
            for tag in all_tags:
                if tag != 'latest' and len(tag) >= 7 and tag.isalnum():
                    return tag
        
        # Fallback to original tag if no git commit found
        return image_uri.split(':')[-1]
    
    except Exception as e:
        print(f"Error getting git commit from ECR: {e}")
        return image_uri.split(':')[-1]

def send_slack_notification(is_migration, old_version_number, new_version_number, old_commit, new_commit, deployment_id, deployment_user, channel, invoked_by):
    is_rollback = new_version_number < old_version_number
    if is_rollback:
        old_version_number, new_version_number = new_version_number, old_version_number
        old_commit, new_commit = new_commit, old_commit
    if new_commit == "latest":
        # Get the image digest for the 'latest' tag
        response = ecr_client.describe_images(
            repositoryName="studio-api",
            imageIds=[{'imageTag': 'latest'}]
        )
        all_tags = response['imageDetails'][0].get('imageTags', [])
        alternate_tags = [tag for tag in all_tags if tag != 'latest']
        if alternate_tags and alternate_tags[0]:
            new_commit = alternate_tags[0]
            
    compare_link = f"<https://github.com/suno-ai/glockenspiel/compare/{old_commit}...{new_commit}|These changes>"
    new_head = f"<https://github.com/suno-ai/glockenspiel/commit/{new_commit}|{new_commit[0:7]}>"
    deploy_link = f"<https://us-east-2.console.aws.amazon.com/codesuite/codedeploy/deployments/{deployment_id}|{deployment_id}>"
    general_message = "After step 2 is completed, old containers are been kept for 15 minutes, please monitor {<#C05BGGY32TU|tech-alerts>} for any issues, and click instant rollback if needed. Migration cannot be rolled back."
    migration_message = ' with migration' if is_migration else ' without migration'
    
    old_version_int = old_version_number.split(':')[-1]
    new_version_int = new_version_number.split(':')[-1]

    actor_verb = f"{deployment_user} is rolling"
    if invoked_by == 'codepipeline.amazonaws.com':
        actor_verb = "CodePipeline is auto-rolling"

    if new_version_int == old_version_int:
        text = f"🚀 {actor_verb} studio api{migration_message} {old_version_int} -> {new_version_int}\nDeploying same head to studio api, new head is {new_head}. Deployment ID: {deploy_link})\n{general_message}"
    elif is_rollback:
        text = f"🚀 {actor_verb} back studio api{migration_message} {old_version_int} -> {new_version_int}\n{compare_link} are been rolled back, new head is {new_head}. Deployment ID: {deploy_link})\n{general_message}"
    else:
        text = f"🚀 {actor_verb} studio api{migration_message} {old_version_int} -> {new_version_int}\n{compare_link} are been deployed, new head is {new_head}. Deployment ID: {deploy_link})\n{general_message}"
        
    message = {
        "channel": channel,
        "icon_emoji": ":cruise_ship:",
        "text": text
    }
    response = requests.post(url, headers=headers, json=message)
    print(response.text)

def lambda_handler(event, context):
    # Print the event details
    print(event)

    # Extract relevant information from the event
    detail = event.get('detail', {})
    user_identity = detail.get('userIdentity', {})
    principal_id = user_identity.get('principalId', 'Unknown')
    user_name = principal_id.split(':')[1]
    print("User name: ", user_name)
    
    invoked_by = user_identity.get('invokedBy')

    # Extract app spec content with backward compatibility
    revision = detail.get('requestParameters', {}).get('revision', {})
    app_spec_content = None
    
    # New format: revision.string.content (YAML string)
    if 'string' in revision and 'content' in revision['string']:
        app_spec_content_str = revision['string']['content']
        print("Using new format: revision.string.content")
        print("App spec content string: ", app_spec_content_str)
        try:
            app_spec_content = parse_codedeploy_appspec_yaml(app_spec_content_str)
        except Exception as e:
            print(f"Error parsing YAML: {e}")
            return {
                'statusCode': 400,
                'body': json.dumps('Failed to parse YAML content')
            }
    # Old format: revision.appSpecContent.content (JSON string)
    elif 'appSpecContent' in revision and 'content' in revision['appSpecContent']:
        app_spec_content_str = revision['appSpecContent']['content']
        print("Using old format: revision.appSpecContent.content")
        print("App spec content string: ", app_spec_content_str)
        try:
            app_spec_content = json.loads(app_spec_content_str)
        except json.JSONDecodeError as e:
            print(f"Error parsing JSON: {e}")
            return {
                'statusCode': 400,
                'body': json.dumps('Failed to parse JSON content')
            }
    else:
        print("Could not find app spec content in either format")
        return {
            'statusCode': 400,
            'body': json.dumps('Could not find app spec content')
        }
    
    print("App spec content: ", app_spec_content)

    response_elements = detail.get('responseElements', {})
    deploymentId = response_elements.get('deploymentId', {})
    print("Deployment ID: ", deploymentId)

    new_task_definition_arn = app_spec_content.get('Resources', {})[0].get('TargetService', {}).get('Properties', {}).get('TaskDefinition', 'Unknown')
    print("New task definition ARN: ", new_task_definition_arn)
    cluster_arn = 'arn:aws:ecs:us-east-2:734185074900:cluster/BackendInfraStack-StudioApiCluster34FACBC9-CTvol2KS9huM' 
    service_arn = 'arn:aws:ecs:us-east-2:734185074900:service/BackendInfraStack-StudioApiCluster34FACBC9-CTvol2KS9huM/StudioApiFargateService'
    if new_task_definition_arn == 'Unknown' or service_arn == 'Unknown' or cluster_arn == 'Unknown':
        print("Unknown task definition ARN or service ARN or cluster ARN, skipping")
        return {
            'statusCode': 200,
            'body': json.dumps('Event processed successfully')
        }

    # Get task definition details
    task_definition = ecs_client.describe_task_definition(
        taskDefinition=new_task_definition_arn
    )
    print("New task_definition: ", task_definition)
    new_image_uri = task_definition['taskDefinition']['containerDefinitions'][0]['image']
    new_commit = get_git_commit_from_ecr_image(new_image_uri)
    print("New commit hash: ", new_commit)

    # Get current task definition of the ECS service
    service = ecs_client.describe_services(
        cluster=cluster_arn,
        services=[service_arn]
    )
    old_task_definition_arn = service['services'][0]['taskDefinition']
    print("Current Task Definition ARN: ", old_task_definition_arn)

    old_version_number = old_task_definition_arn.split('/')[-1]
    new_version_number = new_task_definition_arn.split('/')[-1]

    print("Old version number: ", old_version_number)
    print("New version number: ", new_version_number)


    current_task_definition = ecs_client.describe_task_definition(
        taskDefinition=old_task_definition_arn
    )
    print("Old Task Definition: ", current_task_definition)
    old_image_uri = current_task_definition['taskDefinition']['containerDefinitions'][0]['image']
    old_commit = get_git_commit_from_ecr_image(old_image_uri)
    print("Old commit hash: ", old_commit)
    # Add any additional logic here if needed
    is_migration = False
    hooks = app_spec_content.get('Hooks', {})
    if len(hooks) > 0:
        first_hook = hooks[0]
        if first_hook.get('BeforeInstall') == 'DbMigrationFunctionBeforeInstall':
            is_migration = True

    send_slack_notification(is_migration, old_version_number, new_version_number, old_commit, new_commit, deploymentId, user_name, 'tech-deploy-alerts', invoked_by)
    return {
        'statusCode': 200,
        'body': json.dumps('Event processed successfully')
    }