" Python function call splitter
" Maintainer: Sergey Matveev <stargrave@stargrave.org>
" License: GNU General Public License version 3 of the License or later
"
" This plugin splits Python function call on several lines.
"
"   def foobar(self, foo: str, bar: Some[thing, too]) -> None:
" to
"   def foobar(
"           self,
"           foo: str,
"           bar: Some[thing, too],
"   ) -> None:
"
"   foo(bar, baz)[0]
" to
"   foo(
"       bar,
"       baz,
"   )[0]
"
" You can un-split it using :Undefsplit command on the line where
" splitting starts.
"
" :Defsplit has optional argument specifying how many opening round
" parenthesis must be skipped.
" :Defsplit 1 on foo(baz(baz(...))) produces
"    foo(baz(
"        baz(...),
"    ))
"
" Also there is :Brsplit command behaving similarly, but it splits other
" types of brackets: "{}", "[]".

if exists("g:loaded_defsplit") | finish | endif
let g:loaded_defsplit = 1
if !exists("g:defsplit_shift") | let g:defsplit_shift = "    " | endif

function! s:bracket_find(brs_allowable, line, offset)
    let possible = []
    for bracket in a:brs_allowable
        let found = stridx(a:line, bracket, a:offset)
        if found != -1 | let possible += [found] | endif
    endfor
    return min(possible)
endfunction

function! s:defsplit(brs_allowable, single_line_comma, ...)
    if a:0 == 0 | let skip = 0 | else | let skip = str2nr(a:1) | endif
    let shift = g:defsplit_shift
    let line = getline(".")
    for i in range(len(line))
        if line[i] != g:defsplit_shift[0]
            let prfx = strpart(line, 0, i)
            if line[i : i+3] ==# "def " ||
                \line[i : i+5] ==# "class " ||
                \line[i : i+9] ==# "async def "
                let shift .= shift
            endif
            break
        endif
    endfor
    let brs = {"(": ")", "[": "]", "{": "}"}
    let brfirst = s:bracket_find(a:brs_allowable, line, 0)
    let brlast = strridx(line, brs[line[brfirst]])
    while skip > 0
        let brfirst = s:bracket_find(a:brs_allowable, line, brfirst + 1)
        let brlast = strridx(line, brs[line[brfirst]], brlast - 1)
        let skip -= 1
    endwhile
    let [curly, round, squar, outbuf] = [0, 0, 0, ""]
    let ready = [strpart(line, 0, brfirst + 1)]
    let trailing_comma = 1
    for c in split(line[brfirst + 1 : brlast-1], '\zs')
        if c ==# "*" | let trailing_comma = 0 | endif
        if outbuf ==# "" && c ==# " " | continue | endif
        let outbuf .= c
        if c ==# "," && !curly && !round && !squar
            let ready = add(ready, prfx . shift . outbuf)
            let outbuf = ""
        elseif c ==# "[" | let squar += 1
        elseif c ==# "]" | let squar -= 1
        elseif c ==# "(" | let round += 1
        elseif c ==# ")" | let round -= 1
        elseif c ==# "{" | let curly += 1
        elseif c ==# "}" | let curly -= 1
        endif
    endfor
    if trailing_comma && !(a:single_line_comma == v:true && len(ready) == 1)
        let outbuf = outbuf . ","
    endif
    let ready = add(ready, prfx . shift . outbuf)
    let ready = add(ready, prfx . strpart(line, brlast))
    call append(line("."), ready)
    normal "_dd
endfunction

command! -nargs=? Defsplit call s:defsplit(["("], v:false, <f-args>)
command! -nargs=? Brsplit call s:defsplit(["(", "[", "{"], v:false, <f-args>)
command! -nargs=? Defsplits call s:defsplit(["("], v:true, <f-args>)
command! -nargs=? Brsplits call s:defsplit(["(", "[", "{"], v:true, <f-args>)

command! Undefsplit normal ^v%$J:keepp s/^\(.*\)\([([{]\) \(.*[^,]\),\?\([)\]}]\)\(.*\)$/\1\2\3\4\5<CR>:keepp s/, \?\([)\]}]\+\)$/\1/e<CR>:<CR>