import json
import importlib
import traceback
import sys
import os

import pic16
import pic18
import pic24

running = True
DEBUG = False

ANSI_BLUE = "\x1b[34m"
ANSI_YELLOW = "\x1b[33m"
ANSI_DARK_GREEN = "\x1b[32m"
ANSI_GREEN = "\x1b[1m\x1b[32m"
ANSI_RED = "\x1b[31m"
ANSI_RESET = "\x1b[0m"

COLOR_REGISTER = ANSI_YELLOW
COLOR_FUNCTION = ANSI_GREEN
COLOR_LABEL = ANSI_DARK_GREEN
COLOR_ADDRESS = ANSI_DARK_GREEN

state = {
    'running': True,
    'dirty': False
}

def parse_parts(msg):
    parts = []
    i = 0
    wordstart = 0
    word = ""
    stripped = msg.strip()

    in_quotes = False
    escaping = False

    while i < len(stripped):
        if escaping:
            # escaping means even spaces aren't the end of the word, so
            word = word + stripped[i]
            escaping = False
        else:
            if stripped[i] == ' ':
                if in_quotes:
                    # this isn't actually the end of a word
                    word = word + stripped[i]
                else:
                    parts.append(word)
                    word = ""
            elif stripped[i] == '\\':
                escaping = True
            elif stripped[i] == '"':
                in_quotes = not in_quotes
            else:
                word = word + stripped[i]

        i = i + 1

    parts.append(word)
    word = ""

    # todo: signal unmatched quotes

    return parts

def readnum(s):
    if s.startswith("0x"):
        return int(s[2:], 16)
    else:
        return int(s)

def parse_cmd(cmd):
    cmd_parts = parse_parts(cmd)
    smol_to_big = {
        "o": "open",
        "c": "define-comment",
        "c+": "define-comment",
        "q": "quit",
        "d": "disassemble",
        "h": "help"
    }

    result = {}

    if len(cmd_parts) > 0:
        result['type'] = cmd_parts[0]
        if result['type'] in smol_to_big:
            result['type'] = smol_to_big[result['type']]

        result['params'] = cmd_parts[1:]
        result['raw_text'] = cmd

    if len(cmd_parts) > 1 and cmd_parts[-1][0] == '@':
        result['where'] = resolve_addr(cmd_parts[-1][1:], state)
    elif len(cmd_parts) > 2 and cmd_parts[-2] == '@':
        result['where'] = resolve_addr(cmd_parts[-1], state)

    if 'where' in result and result['where'] is None:
        result['invalid'] = True

    return result

def resolve_addr(where, state):
    # where may be the name of a label or function
    if where in state['notes']['functions-reverse']:
        return state['notes']['functions-reverse'][where]
    elif where in state['notes']['labels-reverse']:
        return state['notes']['labels-reverse'][where]
    else:
        # if it's not there, it better be a number...
        try:
            return readnum(where)
        except ValueError:
            print("{} is not a known function or label.".format(where))

def do_help(cmd, state):
    print("haha, help")

def do_open(cmd, state):
    f = open(cmd['params'][0])
    buf = f.read(6)
    f.close()
    if buf == ":02000":
        return do_openhex(cmd, state)
    else:
        return do_openbin(cmd, state)

def do_openbin(cmd, state):
    return inneropen(cmd, state, lambda x: { 0: x.read() })

def do_openhex(cmd, state):
    return inneropen(cmd, state, lambda x: read_hex(x))

def read_hex(f):
    lines = f.readlines()
    regions = { 0: [] }
    curr_region = regions[0]

    for line in lines:
        if not line.startswith(':'):
            continue
#            raise Exception("invalid hex line, needs to start with a ':', but was: " + line)
        bytecount = int(line[1:3], 16)
        addr = int(line[3:7], 16)
        rec_type = int(line[7:9], 16)
        data_end = bytecount * 2 + 9
        data = line[9:data_end]
        # ignoring checksum because lazy

        if rec_type == 4:
            ext_linear_addr = int(data, 16)
            if ext_linear_addr not in regions:
                regions[ext_linear_addr] = []

            curr_region = regions[ext_linear_addr]
        elif rec_type == 0:
            for i in range(bytecount):
                if len(curr_region) <= addr + i:
                    curr_region.extend([0] * (1 + addr + i - len(curr_region)))
                curr_region[addr + i] = int(data[i*2:i*2 + 2], 16)
        elif rec_type == 1:
            if DEBUG:
                print("Read HEX file ({} regions)".format(len(regions)))
            return regions
        else:
            raise Exception("Unsupported record type: " + str(rec_type))

def do_goto(cmd, state):
    try:
        dest = int(cmd['params'][0])
    except ValueError:
        # parse error, might be a function or label!
        # labels first...
        name = cmd['params'][0]
        if name in state['notes']['labels-reverse']:
            dest = state['notes']['labels-reverse'][name]
        if name in state['notes']['functions-reverse']:
            dest = state['notes']['functions-reverse'][name]
    state['cursor'] = int(cmd['params'][0])

def do_goto_bank(cmd, state):
    param = cmd['params'][0]
    if not param in state['all-regions']:
        state['all-regions'][param] = []

    state['data'] = state['all-regions'][param]

def inneropen(cmd, state, readfn):
    newfile = cmd['params'][0]
    if os.path.isfile(newfile):
        newdata = None
        try:
            f = open(newfile)
            newdata = readfn(f)
            f.close()
        except Exception as e:
            print("Failed to open new file: {}".format(newfile))
            print("Got: {}".format(e))
            print(traceback.format_exc())
            return
        state['notes'] = {
            'comments': {},
            'functions': {},
            'functions-reverse': {},
            'labels': {},
            'labels-reverse': {}
        }
        state['cursor'] = 0
        state['file'] = newfile
        state['all-regions'] = newdata
        state['selected-region'] = newdata.keys()[0]
        state['data'] = newdata[state['selected-region']]
    elif os.path.isdir(newfile):
        print("Cannot open {}, it is a directory".format(newfile))
    else:
        print("File {} does not exist".format(newfile))

    if 'arch-name' not in state or state['arch-name'] is None:
        # default arch to the best cpu, pic16
        do_setarch({ 'params': ['pic16'] }, state)

    if DEBUG:
        print("Opened {} as {}".format(newfile, state['arch-name']))

def do_comment(cmd, state):
    state['notes']['comments'][cmd['where']] = cmd['params'][0]
    state['dirty'] = True

def new_function(name):
    return {
        "name": name,
        "params": None,
        "returns": None
    }

def do_undefine_function(cmd, state):
    name = cmd['params'][0]
    if cmd['params'][0] in state['notes']['functions-reverse']:
        fnaddr = state['notes']['functions-reverse'][cmd['params'][0]]
        del state['notes']['functions'][fnaddr]
        del state['notes']['functions-reverse'][name]
        return
    else:
        print("Function {} is not defined".format(name))

def do_define_function(cmd, state):
    if cmd['params'][0] in state['notes']['functions-reverse']:
        fnaddr = state['notes']['functions-reverse'][cmd['params'][0]]
        if fnaddr != cmd['where']:
            print("Function {} is already defined at {}".format(cmd['params'][0], fnaddr))
            return
        else:
            newname = cmd['params'][0]
            fn = state['notes']['functions'][cmd['where']]
            del state['notes']['functions-reverse'][fn['name']]
            state['notes']['functions-reverse'][newname] = fn
            fn['name'] = newname
    else:
        state['notes']['functions'][cmd['where']] = new_function(cmd['params'][0])
        state['notes']['functions-reverse'][cmd['params'][0]] = cmd['where']
    state['dirty'] = True

def do_define_label(cmd, state):
    state['notes']['labels'][cmd['where']] = cmd['params'][0]
    state['notes']['labels-reverse'][cmd['params'][0]] = cmd['where']
    state['dirty'] = True

def colorize(string, color):
    return "{}{}{}".format(
        color,
        string,
        ANSI_RESET
    )

def do_disassemble(cmd, state):
    if 'data' not in state:
        print("No file currently open")
        return
    if 'cursor' not in state:
        print("No cursor into the file (this is a bug - a file must be open?")
        return

    if 'where' in cmd:
        where = cmd['where']
    else:
        where = state['cursor']

    if len(cmd['params']) > 0:
        count = int(cmd['params'][0])
    else:
        count = 1

    arch = state['arch']

    try:
        disassembled = 0
        while disassembled < count:
            prewhere = where
            (where, instr) = arch.disassemble(state['data'], where)
            if 'ops' in instr:
                newops = instr['ops']
                for i, op in enumerate(instr['ops']):
                    if isinstance(op, dict) and 'type' in op:
                        note = None
                        if op['type'] == "absolutedest":
                            if op['value'] in state['notes']['functions']:
                                note = state['notes']['functions'][op['value']]['name']
                                note = colorize(note, COLOR_FUNCTION)
                            elif op['value'] in state['notes']['labels']:
                                note = state['notes']['labels'][op['value']]
                                note = colorize(note, COLOR_LABEL)
                            else:
                                note = colorize(hex(op['value']), COLOR_ADDRESS)
                        elif op['type'] == "relpostdest":
                            # if there's a thing there replace with it
                            newops[i] = "relpostdest"
                            ea = prewhere + instr['length'] * 2 + op['value']
                            if ea in state['notes']['functions']:
                                note = state['notes']['functions'][ea]['name']
                                note = colorize(note, COLOR_FUNCTION)
                            elif ea in state['notes']['labels']:
                                note = state['notes']['labels'][ea]
                                note = colorize(note, COLOR_LABEL)
                            else:
                                note = "0x{:x} (ip+2{}{})".format(
                                    ea,
                                    "+" if op['value'] >= 0 else "",
                                    hex(op['value'])
                                )
                                note = colorize(note, COLOR_ADDRESS)
                        elif op['type'] == "relpredest":
                            pass
                        elif op['type'] == "register" or op['type'] == "banked-register":
                            nicename = arch.reg_name(op['value'])
                            if nicename is None:
                                if isinstance(op['value'], str):
                                    note = op['value']
                                else:
                                    note = hex(op['value'])
                            else:
                                note = nicename

                            note = colorize(note, COLOR_REGISTER)

                        if not note is None:
                            newops[i] = str(note)

                instr['ops'] = newops

            prefix = None
            if prewhere in state['notes']['functions']:
                fn = state['notes']['functions'][prewhere]
                prefix = colorize(
                    "{:08x}> start of {}\n".format(prewhere, fn['name']),
                    COLOR_FUNCTION
                )
            elif prewhere in state['notes']['labels']:
                label = state['notes']['labels'][prewhere]
                prefix = colorize(
                    "{:08x}: {}\n".format(prewhere, label),
                    COLOR_LABEL
                )
            instrstring = arch.render(instr)
            to_show = "{:08x}: {}".format(prewhere, instrstring)
            if prewhere in state['notes']['comments']:
                to_show = "{: <35} {}; {}{}".format(
                    to_show,
                    ANSI_BLUE,
                    state['notes']['comments'][prewhere],
                    ANSI_RESET
                )
            if not prefix is None:
                to_show = prefix + to_show
            print(to_show)
            disassembled = disassembled + 1
    except Exception as e:
        print("Exception while disassembling: {}".format(e))
        print(traceback.format_exc())

def do_setarch(cmd, state):
    try:
        loaded_arch = importlib.import_module(cmd['params'][0])
    except ImportError:
        print("Cannot find module for arch '{}'".format(cmd['params'][0]))
    except Exception as e:
        print("General error loading '{}': {}".format(cmd['params'][0], e))

    state['arch-name'] = cmd['params'][0]
    state['arch'] = loaded_arch

def dict_from_file(path):
    f = open(path)
    result = json.loads(f.read())
    f.close()
    return result

def dict_to_file(path, data):
    f = open(path, 'w')
    f.write(json.dumps(data, sort_keys=True, indent=2))
    f.close()

# so json.dumps turns numeric keys into strings.
# this breaks some dictionaries.
def fix_keys(obj):
    redo = []
    for k in obj:
        try:
            knum = int(k)
            redo.append(k)
        except ValueError:
            pass

    for k in redo:
        knum = int(k)
        obj[knum] = obj[k]
        del obj[k]

def do_loaddb(cmd, state):
    if len(cmd['params']) == 0:
        if 'file' in state:
            dbpath = state['file'] + '.nrt'
        else:
            print("No db path provided nor file loaded - I don't know what to load!")
            return
    else:
        dbpath = cmd['params'][0]
    if os.path.isdir(dbpath) and os.path.isdir(dbpath + '/.git'):
        if DEBUG:
            print("Loading {} ...".format(dbpath))
        try:
            dbroot = dict_from_file(dbpath + '/root.json')
            do_setarch({ "params": [dbroot['arch-name']] }, state)
            state['notes']['comments'] = dict_from_file(dbpath + '/comments.json')
            fix_keys(state['notes']['comments'])
            state['notes']['functions'] = dict_from_file(dbpath + '/functions.json')
            state['notes']['functions-reverse'] = dict_from_file(dbpath + '/functions-reverse.json')
            fix_keys(state['notes']['functions'])
            state['notes']['labels'] = dict_from_file(dbpath + '/labels.json')
            state['notes']['labels-reverse'] = dict_from_file(dbpath + '/labels-reverse.json')
            fix_keys(state['notes']['labels'])
            state['default-dbpath'] = dbpath
# check db file path matches current file path?
            if DEBUG:
                print("wow you're really using this, huh?")
        except Exception as e:
            print("Error loading db: {}".format(e))
    elif os.path.isdir(dbpath):
        print("Cannot load {}, there is no repository there".format(dbpath))
    elif os.path.isfile(dbpath):
        print("Cannot load {}, it is a file".format(dbpath))
    else:
        print("Cannot load {}, it does not exist".format(dbpath))

def do_savedb(cmd, state):
    if len(cmd['params']) > 0:
        dbpath = cmd['params'][0]
    elif 'default-dbpath' in state:
        dbpath = state['default-dbpath']
    else:
        dbpath = state['file'] + '.nrt'

    if os.path.isfile(dbpath):
        print("File is present, but should be a directory. Cannot save.")
        return
    elif os.path.isdir(dbpath):
        if not os.path.isdir(dbpath + '/.git'):
            print("dbpath is a directory, but there is no git repo there. Opting to not save.")
            return
        else:
            # dbpath exists, and there's a git repo there, we can proceed
            pass
    else:
        # none of the directories exist, so we can start fresh
        os.mkdir(dbpath)
        # TODO: pray the input filename doesn't have a ' in it i guess
        os.system('cd \'{}\' && git init'.format(dbpath))

    to_save = {}
    to_save['arch-name'] = state['arch-name']
    dict_to_file(dbpath + '/root.json', to_save)
    dict_to_file(dbpath + '/comments.json', state['notes']['comments'])
    dict_to_file(dbpath + '/functions.json', state['notes']['functions'])
    dict_to_file(dbpath + '/functions-reverse.json', state['notes']['functions-reverse'])
    dict_to_file(dbpath + '/labels.json', state['notes']['labels'])
    dict_to_file(dbpath + '/labels-reverse.json', state['notes']['labels-reverse'])
    os.system('cd \'{}\' && git add . && git commit -m "automatic save"'.format(dbpath))
    state['dirty'] = False

def do_quit(cmd, state):
    # prompt before saving
    if not 'dirty' in state or state['dirty']:
        do_savedb({ "params": [] }, state)
    state['running'] = False

def do_sh(cmd, state):
    if len(cmd['params']) == 1:
        os.system(cmd['params'][0])
    else:
        print("sh expects exactly one argument (may be a string)")

def do_info(cmd, state):
    if 'file' in state:
        print("File: {}".format(state['file']))
        print("Regions:")
        keys = state['all-regions'].keys()
        keys.sort()
        for key in keys:
            print(" {}{}: {} bytes".format(
                '*' if key == state['selected-region'] else ' ',
                key,
                hex(len(state['all-regions'][key]))
            ))

def readnum(string):
    if string.startswith('0x'):
        return int(string, 16)
    else:
        return int(string)

def do_list_comments(cmd, state):
    print("Comments:")
    for l in state['notes']['comments']:
        print("{}: 0x{:x}".format(state['notes']['comments'][l], l))

def do_list_functions(cmd, state):
    print("Functions:")
    for l in state['notes']['functions']:
        print("{}: 0x{:x}".format(state['notes']['functions'][l]['name'], l))

def do_list_labels(cmd, state):
    print("Labels:")
    for l in state['notes']['labels']:
        print("{}: 0x{:x}".format(state['notes']['labels'][l], l))

def do_hexprint(cmd, state):
    count = readnum(cmd['params'][0])
    start = cmd['where']
    idx = start

    while count > idx - start:
        sys.stdout.write("{:08x}: ".format(idx))
        for i in range(16):
            if count <= idx - start:
                break
            sys.stdout.write("{:02x} ".format(state['data'][idx]),)
            idx = idx + 1
        sys.stdout.write('\n')

cmdmap = {
    "px": do_hexprint,
    "define-comment": do_comment,
    "define-function": do_define_function,
    "undefine-function": do_undefine_function,
    "define-label": do_define_label,
    "list-comments": do_list_comments,
    "list-functions": do_list_functions,
    "list-labels": do_list_labels,
    "open": do_open,
    "openhex": do_openhex,
    "arch": do_setarch,
    "disassemble": do_disassemble,
    "loaddb": do_loaddb,
    "savedb": do_savedb,
    "quit": do_quit,
    "help": do_help,
    "info": do_info,
    "goto": do_goto,
    "goto-bank": do_goto_bank,
    "sh": do_sh
}

def do_cmd_thing(cmd, state):
    if 'invalid' in cmd:
        pass # can't do anything with it.
    elif cmd['type'] in cmdmap:
        cmdmap[cmd['type']](cmd, state)
    else:
        print("I don't recognize the command `{}`".format(cmd['raw_text']))

for i, arg in enumerate(sys.argv):
    # just the path of this script
    # i bet that only applies because i'm running as
    # `python pydare.py ...`
    if i == 0:
        continue
    do_cmd_thing(parse_cmd(arg), state)

# if len(sys.argv) > 1:
#     do_open({"params": [sys.argv[1]]}, state)
#     if len(sys.argv) > 2:
#         do_setarch({"params": [sys.argv[2]]}, state)

import readline
readline.parse_and_bind("")
if os.path.isfile(".pydare_history"):
    readline.read_history_file(".pydare_history")

while state['running']:
    try:
        do_cmd_thing(parse_cmd(raw_input("> ")), state)
    except Exception as e:
        print("Unhandled exception: {}".format(e))
        print(traceback.format_exc())

if not os.path.isdir(".pydare_history"):
    readline.write_history_file(".pydare_history")