function! AlignWithTopLine() range
    " Handle case when selection is at the start of the file
    if a:firstline <= 1
        " If at the start of file, use default indentation (0)
        let reference_indent = 0
        let above_line = ""
    else
        " Find the first non-empty line above the selection
        let above_line_num = a:firstline - 1
        while above_line_num >= 1 && getline(above_line_num) =~ '^\s*$'
            let above_line_num = above_line_num - 1
        endwhile
        
        " If we couldn't find a non-empty line, use default indentation
        if above_line_num < 1
            let reference_indent = 0
            let above_line = ""
        else
            let above_line = getline(above_line_num)
            let reference_indent = indent(above_line_num)
        endif
    endif
    
    " Check characteristics of the line above (if it exists)
    let ends_with_colon = !empty(above_line) && above_line =~# ':$'
    let starts_with_special = !empty(above_line) && above_line =~# '^\s*\(with\|await\|if\|for\|while\|def\|class\)'
    let is_continuation = !empty(above_line) && above_line =~# '\((\|\[\|{\).*\(,\|\\\)$'
    
    " Determine target indentation
    let target_indent = reference_indent
    
    " Handle different cases for alignment
    if ends_with_colon || starts_with_special
        " For block starts or special constructs, add one level of indentation
        let target_indent = reference_indent + &shiftwidth
    elseif is_continuation
        " For line continuations (ending with comma or backslash), align with extra indent
        let target_indent = reference_indent + &shiftwidth
    endif
    
    " Get the first line's current indentation as reference for relative shifts
    let first_line = getline(a:firstline)
    let first_line_indent = indent(a:firstline)
    let first_line_content = substitute(first_line, '^\s*', '', '')
    
    " Check if the first line of selection is a continuation of a previous statement
    let first_line_is_continuation = !empty(above_line) && 
                                   \ (above_line =~# '\((\|\[\|{\|\\$\|,$\)' && 
                                   \ first_line_content !~# '^\()\|]\|}\)')
    
    " Calculate the indentation shift needed
    let indent_shift = 0
    if first_line_is_continuation
        let indent_shift = target_indent - first_line_indent
    else
        let indent_shift = target_indent - first_line_indent
    endif
    
    " Apply the indentation to all lines in the selection
    for lineno in range(a:firstline, a:lastline)
        let line = getline(lineno)
        let current_indent = indent(lineno)
        let current_line_content = substitute(line, '^\s*', '', '')
        
        " Skip empty lines
        if current_line_content == ''
            continue
        endif
        
        " Calculate new indentation preserving relative structure
        let new_indent = current_indent + indent_shift
        
        " Special handling for closing brackets/braces/parentheses
        if current_line_content =~# '^\()\|]\|}\)'
            " Closing brackets align with the opening line (one level back)
            let new_indent = max([0, target_indent - &shiftwidth])
        endif
        
        " Ensure indentation is never negative
        let new_indent = max([0, new_indent])
        
        " Create the new line with proper indentation
        let new_line = repeat(' ', new_indent) . current_line_content
        
        call setline(lineno, new_line)
    endfor
endfunction

function! ConditionalAlign()
    if &filetype == 'python'
        " Store the original selection boundaries
        let l:start_line = a:firstline
        let l:end_line = a:lastline
        
        " Apply indent to the selection. autopep8 will not align if 
        " with xx: 
        " dosomethin 
        " if there are not indentation 
        norm gv4>
        
        " Run autopep8 on the selection, assume indentation = 0 
        silent execute l:start_line . ',' . l:end_line . '!autopep8 -'
        " Re-indent to above line
        
        silent execute l:start_line . ',' . l:end_line . 'call AlignWithTopLine()'
    else
        normal! gv=
    endif
endfunction

vnoremap = :'<,'>call ConditionalAlign()<CR>