#!/usr/bin/env python3

#  how to detect and colors:
#    files changed in index
#    changes not staged for commit
#    untracked files

import sys
import sh
from visidata import *

__version__ = 'saul.pw/vgit v0.1'

option('color_git_unmodified', 'yellow')
option('color_git_add_staged', 'green')
option('color_git_wt_new', 'cyan')
option('color_git_mod_unstaged', 'brown')
option('vgit_show_ignored', True)
option('diff_algorithm', 'minimal')

command('x', 'i = input("git ", type="git"); vd.push(TextSheet(i, sh.git(*i.split()).split("\\n")))', 'execute arbitrary git command')
command('B', 'vd.push(GitBranches())', 'execute arbitrary git command')
command('O', 'vd.push(GitOptions())', 'push sheet of git options')

def open_git(p):
    return GitStatus(p)

class GitSheet(Sheet):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.extra_args = []
        self.branch = ''

    def git(self, *args, **kwargs):
        args = list(args) + self.extra_args
        self.extra_args.clear()
        status(sh.git(*args, **kwargs))
        self.reload()

    def leftStatus(self):
        return ('‹%s› ' % self.branch) + super().leftStatus()

    def reload(self):
        self.branch = sh.git('rev-parse', '--abbrev-ref', 'HEAD').strip()


class GitBranches(GitSheet):
    def __init__(self):
        super().__init__('branches')
        self.columns = [
            Column('branch')
        ]

        self.command('a', 'git("checkout", "-b", input("create branch: ", type="branch"))', 'create a new branch off the current checkout')
        self.command('d', 'git("branch", "--delete", cursorRow)', 'delete this branch')
        self.command('e', 'git("branch", "-v", "--move", cursorRow, editCell(0))', 'rename this branch')
        self.command([ENTER, 'c'], 'git("checkout", cursorRow)', 'checkout this branch')
        self.command('m', 'git("merge", cursorRow)', 'merge this branch into the current branch')

    def reload(self):
        super().reload()
        self.rows = []
        for line in sh.git('--no-pager', 'branch', '--no-color',
#                '--all',  # local and remote branches
                '--list').splitlines():
            self.rows.append(line[2:])



class GitStatus(GitSheet):
    def __init__(self, p):
        super().__init__(p.relpath(''), p)

        self.columns = [Column('path', getter=lambda r,s=self: s.filename(r) + (r.is_dir() and '/' or ''), cache=True),
                      Column('status', getter=lambda r,s=self: s.statusText(s.git_status(r)), width=8),
                      Column('staged', getter=lambda r,s=self: s.git_status(r)[2]),
                      Column('unstaged', getter=lambda r,s=self: s.git_status(r)[1]),
                      Column('type', getter=lambda r: r.is_dir() and '/' or r.suffix, width=0),
                      Column('size', type=int, getter=lambda r: r.filesize),
                      Column('mtime', type=date, getter=lambda r: r.stat().st_mtime),
                    ]

        self.addColorizer('row', 3, GitStatus.rowColor)

        self.command('f', 'sheet.extra_args.append("-f")', 'force next action with "-f"')

        self.command('a', 'git("add", filename(cursorRow))', 'add this new file or modified file to staging')
        self.command('m', 'git("mv", filename(cursorRow), input("rename file to: ", value=filename(cursorRow)))', 'rename this file')
        self.command('d', 'git("rm", filename(cursorRow))', 'stage this file for deletion')
        self.command('w', 'git("reset", "HEAD", filename(cursorRow))', 'unstage this file')
        self.command('c', 'git("checkout", filename(cursorRow))', 'checkout this file')
        self.command('ga', 'git("add", *[filename(r) for r in selectedRows])', 'add all selected files to staging')
        self.command('gd', 'git("rm", *[filename(r) for r in selectedRows])', 'delete all selected files')
        self.command(['^S', 'C'], 'git("commit", "-m", input("commit message: "))', 'commit changes')
        self.command('V', 'vd.push(TextSheet(filename(cursorRow), Path(filename(cursorRow))))', 'open file')
        self.command('i', 'open(workdir+"/.gitignore", "a").write(filename(cursorRow)+"\\n")', 'add file to toplevel .gitignore')
        self.command('gi', 'open(workdir+"/.gitignore", "a").write(input("add wildcard to .gitignore: "))', 'add input line to toplevel .gitignore')  # path, filename

        self.command(ENTER, 'vd.push(HunksSheet(cursorRow))', 'push sheet of diffs for this file')

        self.command('g/', 'vd.push(GitGrep(input("git grep: ")))', 'find in all files')
#
        self._cachedStatus = {}

    def statusText(self, st):
        vmod = {'A': 'add', 'D': 'rm', 'M': 'mod', 'T': 'chmod', '?': 'out', '!': 'ignored'}
        x, y = st[0]
        if st == '??': # untracked
            return 'new'
        elif st == '!!':  # ignored
            return 'ignored'
        elif x != ' ' and y == ' ': # staged
            ret = vmod.get(x, '!') + ' '
            return ret
        elif y != ' ': # unstaged
            ret = vmod.get(y, '!')
            return ret
        else:
            return ''

    @property
    def workdir(self):
        return str(self.source)

    @staticmethod  # due to the way colorizers are called
    def rowColor(self, c, row, v):
        st = self.git_status(row)[0]
        x, y = st
        if st == '??': # untracked
            return 'yellow'
        elif st == '!!':  # ignored
            return '237 blue'
        elif x in 'AMD' and y == ' ': # staged add/mod/rm
            return 'green'
        elif y in 'AMD': # unstaged add/mod/rm
            return 'red'

        return 'white'

    def git_status(self, r):
        return self._cachedStatus.get(self.filename(r), ["FU", None, None])

    def ignored(self, fn):
        return not options.vgit_show_ignored and self._cachedStatus[fn][0] == '!!'

    def filename(self, r):
        return r.relpath(str(self.source) + '/')

    def reload(self):
        super().reload()
        filenames = dict((self.filename(p), Path(self.filename(p))) for p in self.source.iterdir() if p.name not in ('.git', '..'))  # files in working dir

        self.rows = []
        self._cachedStatus.clear()
        for fn in sh.git('ls-files', '-z').split('\0'):
            self._cachedStatus[fn] = ['  ', None, None]  # status, adds, dels

        for status_line in sh.git('status', '-z', '-unormal', '--ignored').split('\0'):
            if status_line:
                st = status_line[:2]
                fn = self.filename(Path(status_line[3:]))
                if fn not in filenames:
                    if not self.ignored(fn):
                        self.rows.append(Path(fn))
                self._cachedStatus[fn] = [st, None, None]

        for line in sh.git('diff-files', '--numstat', '-z').split('\0'):
            if not line: continue
            adds, dels, fn = line.split('\t')
            cs = self._cachedStatus[fn]
            cs[1] = '+%s/-%s' % (adds, dels)

        for line in sh.git('diff-index', '--cached', '--numstat', '-z', 'HEAD').split('\0'):
            if not line: continue
            adds, dels, fn = line.split('\t')
            cs = self._cachedStatus[fn]
            cs[2] = '+%s/-%s' % (adds, dels)

        self.rows.extend([p for k, p in filenames.items() if not self.ignored(k)])

        self.recalc()
        self.rows.sort(key=lambda r,col=self.columns[-1]: col.getValue(r), reverse=True)


class HunksSheet(GitSheet):
    def __init__(self, *files):
        super().__init__('hunks', *files)
        self.columns = [
            ColumnItem('origfn', 0, width=0),
            ColumnItem('filename', 1),
            ColumnItem('context', 2),
            ColumnItem('leftlinenum', 3),
            ColumnItem('leftcount', 4),
            ColumnItem('rightlinenum', 5),
            ColumnItem('rightcount', 6),
        ]

        self.command(ENTER, 'vd.push(HunkViewer(sheet, cursorRow))', 'view the diff for this hunk')
        self.command('g^J', 'vd.push(HunkViewer(sheet, *(selectedRows or rows)))', 'view the diffs for the selected hunks')
        self.command('V', 'vd.push(TextSheet("diff", "\\n".join(cursorRow[7])))', 'apply this hunk to the index')
        self.command('a', 'apply(cursorRow)', 'apply this hunk to the index')

    def apply(self, hunk):
        self.git("apply", "--cached", "-p0", "-", _in="\n".join(hunk[7]) + "\n")
        status('applied hunk (lines %s-%s)' % (hunk[3], hunk[3]+hunk[4]))

    def reload(self):
        super().reload()
        diffout = sh.git('--no-pager', 'diff',
                '--diff-algorithm=' + options.diff_algorithm,
                '--patch',
                '--no-color',
                '--no-prefix', *[p.resolve() for p in self.sources])
        self.rows = []
        leftfn = ''
        rightfn = ''
        header_lines = None
        for line in diffout.splitlines():
            if line.startswith('---'):
                header_lines = [line]  # new file
                leftfn = line[4:]
            elif line.startswith('+++'):
                header_lines.append(line)
                rightfn = line[4:]
            elif line.startswith('@@'):
                header_lines.append(line)
                _, linenums, context = line.split('@@')
                leftlinenums, rightlinenums = linenums.split()
                leftstart, leftcount = leftlinenums[1:].split(',')
                rightstart, rightcount = rightlinenums[1:].split(',')
                self.rows.append((leftfn, rightfn, context, int(leftstart), int(leftcount), int(rightstart), int(rightcount), header_lines))
                header_lines = header_lines[:2]  # keep file context
            elif line[0] in ' +-':
                self.rows[-1][-1].append(line)


class HunkViewer(GitSheet):
    def __init__(self, srchunks, *hunks):
        super().__init__('hunk', *hunks)
        self.srchunks = srchunks
        self.columns = [
            ColumnItem('left', 1, width=vd().windowWidth//2-1),
            ColumnItem('right', 2, width=vd().windowWidth//2-1),
        ]
        self.addColorizer('row', 4, HunkViewer.colorDiffRow)
        self.command('a', 'srchunks.apply(sources.pop(0)); reload()', 'apply this hunk to the index and move to the next hunk')
        self.command(ENTER, 'sources.pop(0); reload()', 'move to the next hunk without applying this hunk')
        self.command('d', 'source[7].pop(cursorRow[3]); reload()', 'delete a line from the patch')

    def reload(self):
        super().reload()
        if not self.sources:
            self.vd.remove(self)
            return

        fn, _, context, linenum, _, _, _, patchlines = self.source
        self.name = '%s:%s' % (fn, linenum)
        self.rows = []
        nextDelIdx = None
        for line in patchlines[3:]:  # diff without the patch headers
            typech = line[0]
            line = line[1:]
            if typech == '-':
                self.rows.append([typech, line, None])
                if nextDelIdx is None:
                    nextDelIdx = len(self.rows)-1
            elif typech == '+':
                if nextDelIdx is not None:
                    if nextDelIdx < len(self.rows):
                        self.rows[nextDelIdx][2] = line
                        nextDelIdx += 1
                        continue

                self.rows.append([typech, None, line])
                nextDelIdx = None
            elif typech == ' ':
                self.rows.append([typech, line, line])
                nextDelIdx = None
            else:
                continue  # header

    def colorDiffRow(self, c, row, v):
        if row[1] != row[2]:
            if row[1] is None:
                return 'green'  # addition
            elif row[2] is None:
                return 'red'  # deletion
            else:
                return 'yellow'  # difference


class GitGrep(GitSheet):
    def __init__(self, regex):
        super().__init__(regex, regex, columns=[
            ColumnItem('filename', 0),
            ColumnItem('linenum', 1),
            ColumnItem('line', 2),
        ])
        self.command(ENTER, 'vd.push(openSource(cursorRow[0], filetype="txt")).cursorRowIndex = int(cursorRow[1])-1', 'go to this match')

    def reload(self):
        self.rows = []
        for line in sh.git('--no-pager', 'grep', '--no-color', '-z',
                '--line-number',
                '--ignore-case',
                self.source).splitlines():
            self.rows.append((line.split('\0')))


class GitOptions(GitSheet):
    CONFIG_CONTEXTS = ('local', 'local', 'global', 'system')
    def __init__(self):
        super().__init__('git config')
        self.columns = [Column('option', getter=lambda r: r[0])]
        for i, ctx in enumerate(self.CONFIG_CONTEXTS[1:]):
            self.columns.append(Column(ctx, getter=lambda r, i=i: r[1][i], setter=self.config_setter(ctx)))

        self.nKeys = 1

        self.command('d', 'git("config", "--unset", "--"+CONFIG_CONTEXTS[cursorColIndex], cursorRow[0])', 'unset this config value')
        self.command('gd', 'for r in selectedRows: git("config", "--unset", "--"+CONFIG_CONTEXTS[cursorColIndex], r[0])', 'unset this config value')
        self.command('e', 'i=(cursorVisibleColIndex or 1); visibleCols[i].setValues([cursorRow], editCell(i)); sheet.cursorRowIndex += 1', 'edit this option')
        self.command('ge', 'i=(cursorVisibleColIndex or 1); visibleCols[i].setValues(selectedRows, input("set selected to: ", value=cursorValue))', 'edit this option for all selected rows')
        self.command('a', 'git("config", "--add", "--"+CONFIG_CONTEXTS[cursorColIndex], input("option to add: "), "added")', 'add new option')

    def config_setter(self, ctx):
        def setter(r, v):
            self.git('config', '--'+ctx, r[0], v)
        return setter

    def reload(self):
        super().reload()
        opts = {}
        for i, ctx in enumerate(self.CONFIG_CONTEXTS[1:]):
            try:
                for line in sh.git('config', '--list', '--'+ctx, '-z').split('\0'):
                    if line:
                        k, v = line.splitlines()
                        if k not in opts:
                            opts[k] = [None, None, None]
                        opts[k][i] = v
            except:
                pass # exceptionCaught()

        self.rows = sorted(list(opts.items()))

vd().status(__version__)

setGlobal('sh', sh)
setGlobal('HunksSheet', HunksSheet)
setGlobal('HunkViewer', HunkViewer)
setGlobal('GitBranches', GitBranches)
setGlobal('GitGrep', GitGrep)
setGlobal('GitOptions', GitOptions)
fn = sys.argv[1] if sys.argv[1:] else '.'
os.chdir(fn)
options.textwrap = False
run([GitStatus(Path('.'))])
