import atexit
import os
import argparse
from pwndbg.commands.context import contextoutput
from pwnlib.constants import linux
linux.riscv64 = None

class Tmux(gdb.Command):
    pane_args = '-P -F "#{pane_id}:#{pane_tty}" -d "cat -"'
    terminal = 'alacritty -o font.size=10 -e'

    def __init__(self):
        super().__init__('tmux', gdb.COMMAND_USER)
        self.panes = {}
        self.parser = argparse.ArgumentParser('tmux')
        self.parser.add_argument('layout', type=int, nargs='?')
        self.layouts = [self.layout_single, self.layout_double]

    def reset(self):
        for s in ('stack', 'disasm', 'regs', 'backtrace', 'legend'):
            contextoutput(s, None, clearing=False)
        gdb.execute(f'tty')
        gdb.execute('set context-stack-lines 8')

    def __getattr__(self, name):
        return self.panes[name]

    def splitw(self, *args):
        return self.tmux(f'splitw {' '.join(str(x) for x in args)} {self.pane_args}').read().strip().split(':')

    def session(self, name):
        return self.tmux(f'new-session -s {name} {self.pane_args}').read().strip().split(':')

    def tmux(self, command):
        return os.popen(f'tmux {command}' if isinstance(command, str) else f'tmux ' + '\\; '.join(command))

    def invoke(self, args, from_tty):
        try:
            args = self.parser.parse_args(args.split())
        except:
            return
        if self.panes:
            self.kill_panes()
            self.reset()
            if args.layout is None:
                return
        if args.layout is None:
            args.layout = 0
        else:
            args.layout -= 1
        if 0 <= args.layout < len(self.layouts):
            gdb.execute('set context-stack-lines 40')
            self.layouts[args.layout]()
            self.set_context()
            atexit.register(self.kill_panes)
        else:
            print('no such layout')

    def kill_panes(self):
        for p in self.panes.values():
            self.tmux(f'kill-pane -t {p[0]}').read()
        self.panes = {}

    def set_pane(self, name, pane):
        self.panes[name] = pane

    def layout_single(self):
        self.set_pane('stack', self.splitw())
        self.set_pane('backtrace', self.splitw('-h', '-t', self.stack[0]))
        self.set_pane('regs', self.splitw('-h', '-t', self.backtrace[0]))
        self.set_pane('output', self.splitw('-t', self.backtrace[0]))
        self.set_pane('disasm', self.splitw('-h'))

    def layout_double(self):
        session = f'gdb_{os.getpid()}'
        self.set_pane('disasm', self.session(session))
        self.set_pane('stack', self.splitw('-h', '-t', self.disasm[0]))
        self.set_pane('regs', self.splitw('-t', self.disasm[0]))
        self.set_pane('backtrace', self.splitw('-t', self.regs[0], '-l', 1))
        self.set_pane('output', self.splitw('-h', '-b'))
        os.popen(f'{self.terminal} tmux attach -t {session}')

    def set_context(self):
        gdb.execute(f'tty {self.output[1]}')
        self.tmux(f'select-pane -t {p[0]} -T {sec}' for sec, p in self.panes.items())
        for s, p in self.panes.items():
            contextoutput(s, p[1], clearing=True, banner=None)
        contextoutput("legend", self.disasm[1], clearing=False)

Tmux()