
from staq.archs import X86_cdecl, Arm32
from staq.stack import Stack, StackFrame, StackCell
from staq.function import Function, FunctionArgument, parseCFile
from staq.helper import splitWithEscapes
import re
import yaml
import json
import argparse
import re

import importlib.resources as pkg_resources
import builtins
from cliify import commandParser, command, splitWithEscapes


sizeDict = {
    "void": 0,
    "char": 1,
    "unsigned char": 1,
    "signed char": 1,
    "short": 2,
    "unsigned short": 2,
    "int": 4,
    "unsigned int": 4,
    "long": 4,
    "unsigned long": 4,
    "long long": 8,
    "unsigned long long": 8,
    "float": 4,
    "double": 8,
    "long double": 16,  # This can vary, but 16 bytes is a common size
    "int8": 1,
    "uint8": 1,
    "int16": 2,
    "uint16": 2,
    "int32": 4,
    "uint32": 4,
    "int64": 8,
    "uint64": 8,
    "va_list": 4,
}


def nullPrint(*args, **kwargs):
    pass

builtins.print = nullPrint

@commandParser( subparsers=['arch'],collapsed=True)
class StackSession():
    def __init__(self, arch = None):
        self.stack = Stack()
        self.arch = None
        #self.arch = X86_cdecl(session=self, stack=self.stack)
        self.functions = {}
        self.parser = argparse.ArgumentParser(description='Stack Visualizer', add_help=False)
        self.subparsers = None
        self.loadLibcFuntions()
        self.commands = []
        self.update = False
        self.status = ""
        self.cmdIdx = 0

    def setArch(self, arch):
        if arch == 'x86':
            self.arch = X86_cdecl(session=self, stack=self.stack)
        elif arch == 'arm':
            self.arch = Arm32(session=self, stack=self.stack)
        else:
            self.arch = X86_cdecl(session=self, stack=self.stack)

    def addHistory(self,cmd):

        if self.cmdIdx < len(self.commands):
            self.commands = self.commands[:self.cmdIdx]

        self.commands.append(cmd)
        self.cmdIdx+=1

    def stepForward(self):

        if self.cmdIdx < len(self.commands):
            self.parseCommand(self.commands[self.cmdIdx], addHistory=False)
            self.cmdIdx +=1

    def stepBack(self):

        if self.cmdIdx > 0:

            self.cmdIdx -= 1
            self.stack.clear()
            
            for i in range(0,self.cmdIdx):
                self.parseCommand(self.commands[i], addHistory=False)

    def print(self,text):
        print(text)

    def refreshOutput(self):

        ansi = self.stack.toAnsi( showAddress=True)
        print(ansi)

    def loadYaml(self, obj):


        if isinstance(obj, str):
            with open(obj) as f:
                obj = yaml.safe_load(f)


        order = 'normal'

        if 'stackBase' in obj:
            self.stack.baseAddress = obj['stackBase']
        
        if 'order' in obj:
            order = obj['order']

        if 'functions' in obj:
            for func in obj['functions']:
                self.addFunction(Function.fromString(func))
    

        if 'stack' in obj:

            nodes = []
            if order == 'normal':
                nodes = reversed(obj['stack'])
            else:
                nodes = obj['stack']

            for node in nodes:
                
               if isinstance(node, dict) and 'function' in node:
                   #Create a new frame 
                   frame = StackFrame.fromObj(node, order= order)
                   self.stack.pushFrame(frame)
                   self.stack.currentFrame = None
               else:
                    cell = StackCell.fromObj(node)
                    self.stack.push(cell)
        
        self.stack.applyAddresses()


    
    def getCurrentFunction(self):

        if self.stack.currentFrame:

            if self.stack.currentFrame.function:
                return self.stack.currentFrame.function
            
        
        return None

    def addFunction(self, function):
        self.functions[function.name] = function


    def loadLibcFuntions(self):
        
        data_path = pkg_resources.files('staq.data').joinpath('libc.c')

        with open(data_path) as f:
            functions = parseCFile(f.read())
            for f in functions:
                self.addFunction(f)

        # for f in libc_functions:
        #     newFunc = Function.fromString(f)

        #     self.addFunction(newFunc)

   
    @command
    def frame(self, name, color = None):
        """
            frame <name> --color <color>
        """
        frame = StackFrame(name)
        if color:
            frame.color = color
        self.stack.pushFrame(frame)



    def parseFunction(self, line):

        function = Function.fromString(line)
        self.addFunction(function)


    # def functionCmd(self,args):
    #     line = " ".join(args.function)
    #     function = Function.fromString(line)
    #     self.addFunction(function)
    
    @command( completions={'file': ['$file']} )
    def save(self, file = None):

        filename = file




        if filename.endswith('.html'):
            with open(filename, "w") as f:
                f.write(self.stack.toHtml())

        elif filename.endswith('.png'):

            try:
                self.stack.generatePng(filename)
            except Exception as e:
                self.print("Failed to generate image. This is likely due to not having 'chrome' installed which is required for html2image. Please install chrome and try again.")

        else:
            with open(filename, "w") as f:
                for cmd in self.commands:
                    if not cmd.startswith("save"):
                        if not cmd.endswith("\n"):
                            cmd += "\n"
                        f.write(cmd)


                
        

    @command( completions={'file': ['$file']} )
    def load(self, file: str = None):
        filename = file
        self.update = False

        if filename.endswith('.c'):

            with open(filename) as f:
                functions = parseCFile(f.read())

            self.status = f"Loaded {len(functions)} functions from {filename}"

            for func in functions:
                self.addFunction(func)

    def getCellRefs(self) -> list:
        """
        Get a list of all cell references in the stack
        """
        cells = self.stack.getCells()
        refs = []

        for cell in cells:
            refs.append(f"&{cell.label}")

        return refs

    @command( completions={'file': ['$file'], 'data': lambda  self: self.getCellRefs(), 'address': lambda  self: self.getCellRefs(),  } )
    def write(self, data = None, address = None, file = None ):
        """
            write [data1,data2,data3] , address: &cell+4
            write &cell+16: [data1,data2,data3]
        
        """
        

        if file:
            #read file as bytes 
            with open(file, "rb") as f:
                data = f.read()
                data = list(data)


        if isinstance(data, str):
            if not data.startswith("[") and not data.endswith("]"):
                data = f"[{data}]"
            data = yaml.safe_load(data) #use yaml to parse the list so we get the correct types/escapes

        self.stack.write(data, address)

    @command( completions={'address': lambda  self: self.getCellRefs(), 'note': lambda  self: self.getCellRefs(), 'color': ['red', 'green', 'blue', 'yellow', 'purple'] } )
    def note(self, note = None, address = None, color = 'red'):
        """
            note 0x1234: note on local variable
            note 0x1234: note on local variable --color red
        """

        self.stack.setNote(note=note, address=address, color=color)

    @command(completions={'function': lambda self: self.functions.keys()})
    def run(self, function, limit = 20):

        called = True

        functionCall = None
        
        if function:
            function = function.strip()
            if function == "":
                function = None

        if function:
            functionCall = function
        elif self.stack.currentFrame and self.stack.currentFrame.function:
            function = self.stack.currentFrame.function

            if len(function.calls) > 0:
                functionCall  = function.calls[0]

        if not functionCall:
            functionCall = "main()"

        self.arch.call(functionCall)
        limit -= 1

        while  called and (limit > 0):
            
            if self.stack.currentFrame and self.stack.currentFrame.function: 
                function = self.stack.currentFrame.function

                if len(function.calls) > 0:
                    call = function.calls[0]
                    self.arch.call(call)
                    called = True

            limit -= 1


    @command
    def clear(self):
        self.stack.clear()


    def tryParseCommand(self, line, addHistory = True):
        """
        Try to parse a command. If it fails, try to parse with calling convention
        """
        plainCmd = line.split(" ")[0]

        if addHistory and plainCmd not in ['save','s', 'load','l']:
            self.addHistory(line)


        cmds = splitWithEscapes(line,";")

        for cmd in cmds:
            words = cmd.split(" ")
            firstWord = words[0]

            #If first word is a var declaration for a pointer, move the * to the beginning of the name 
            # i.e int* foo = 0x1234 becomes int *foo = 0x1234
            # This is to allow the command parser to work with the declaration
            if firstWord.endswith("*"):
                remaining = cmd[len(firstWord):].lstrip()
                firstWord = firstWord[:-1]
                cmd = firstWord + " *" + remaining
            
            if firstWord in ["write","note"]:
                """
                    if using shorthand for write, i.e. write &cell+16: [data1,data2,data3]
                """

                if (words[1].startswith("&") or words[1].startswith("0x")) and words[1].endswith(":"):
                    #write &cell+16: [data1,data2,data3]
                    cmd = f"{firstWord} {" ".join(words[2:])} , address: {words[1]}"
                    

                

            self.parseCommand(cmd.strip())
            self.refreshOutput()





    def showHelp(self, command):
            
            if command in self.subparsers.choices:
                self.print(self.subparsers.choices[command].format_help() + "\n\n [Enter] to continue...")
            else:
                self.print(self.parser.format_help() + "\n\n [Enter] to continue...")

            #self.print("[Enterr] to continue...")

    def tryParseLocalVar(self, line):
        
        parts = line.split("=")

        
        revar = re.compile(r"(\w[\w\s\*]*\w)\s+(\w+)(\[(\d*)\])?")

        ptr = False

        decl = parts[0]
        if "*" in decl:
            ptr = True
            decl = decl.replace("*","")

        match = revar.match(decl)
        val = []
        if len(parts) > 1:
            parts[1] = parts[1].split("--")[0]
            val = parts[1].replace(";","")
            val = yaml.safe_load(val)

        if match:
            varType = match.group(1).strip()

            varName = match.group(2)

            arraySize = match.group(4) if match.group(4) else None

            cell = StackCell(varName)

            if "_t" in varType:
                varType = varType.replace("_t","")

            if varType not in sizeDict:
                return None

            
            if ptr:
                sizeBytes = 4
            else:
                sizeBytes = sizeDict.get(varType, 4)

            if arraySize:
                sizeBytes = sizeBytes * int(arraySize)
            
            size = int((sizeBytes + 3) / 4)

            cell.setWords(val)
            cell.setSize(size)
            cell.label = varName


            line = line.replace(match.group(0), 'local')


            return cell

        else:
            return None

