
from staq.conventions import CdeclConvention
from staq.stack import Stack, StackFrame, StackCell
from staq.function import Function, FunctionArgument, parseCFile
import re
import yaml
import json
import argparse

import importlib.resources as pkg_resources
import builtins


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

class StackSession():
    def __init__(self):
        self.stack = Stack()
        self.conv = CdeclConvention(session=self, stack = self.stack)
        self.functions = {}
        self.parser = argparse.ArgumentParser(description='Stack Visualizer', add_help=False)
        self.subparsers = None
        self.init_args()
        self.loadLibcFuntions()
        self.commands = []
        self.update = False
        self.status = ""

    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)

    def parseFrame(self,words):
        
        line = " ".join(words[1:])
        self.stack.pushFrame(StackFrame(line))

    def frameCmd(self, args):
        frame = StackFrame(args.name)
        if args.color:
            frame.color = args.color
        self.stack.pushFrame(frame)

    def popCmd(self, args):

    
        if args.count and isinstance(args.count,str):
            if args.count.lower() == 'frame':
                self.stack.popFrame()
            elif args.count.lower().replace("$","") in self.conv.registers:

                word = self.stack.pop()
                self.conv.setReg(args.count.lower().replace("$",""), word.value )
            else:
                count = int(args.count)
                self.stack.pop(count)
        else:
            self.stack.pop()

    def retCmd(self, args):

        self.conv.ret()

    def jmpCmd(self, args):

        self.conv.jmp(args.address)

    def pushCmd(self, args):
        cell = StackCell()

        if args.value:

            line = " ".join(args.value)

            parts = line.split(":")

            if len(parts) > 1:
                cell.label = parts[0]

                words = parts[1].strip().split(",")
                words = [x.strip() for x in words]

                cell.setWords(words)
            else:
                words = parts[0].split(",")
                words = [x.strip() for x in words]
                cell.setWords(parts[0].split(","))


        if args.size:
            cell.setSize(args.size)
        
        if args.label:
            cell.label = args.label

        if args.address:
            cell.address = args.address

        if args.note:

            note = " ".join(args.note)

            cell.setNote(note)



        self.stack.push(cell)


    def parseFunction(self, line):

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

   
    def callCmd(self, args):
            

            line = " ".join(args.function)
    
            self.conv.call(line)

    def leaveCmd(self, args):
        self.conv.leave()

    def functionCmd(self,args):
        line = " ".join(args.function)
        function = Function.fromString(line)
        self.addFunction(function)
    
    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
            
    def localCmd(self, cell, args):

        if args.note:
            cell.setNote(" ".join(args.note))

        if args.color:
            cell.color = args.color

        self.stack.push(cell)

    def saveCmd(self, args):

        filename = args.filename




        if filename.endswith('.stackinit'):
            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)
        
        elif filename.endswith('.html'):
            with open(filename, "w") as f:
                f.write(self.stack.toHtml())

        elif filename.endswith('.png'):
            self.stack.generatePng(filename)


                
        


    def loadCmd(self, args):
        filename = args.filename
        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 writeCmd(self, args):

        valLine = " ".join(args.values)
        values = valLine.split(",")

        #strip 
        values = [x.strip() for x in values]

        address  = None

        if args.address:
            address = args.address

        self.stack.write(values, address)


    def noteCmd(self, args):
        note = " ".join(args.note)
        self.stack.setNote(note, args.address, args.color)

    def run(self, callLine = None, limit = 20):

        called = True

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

        if callLine:
            functionCall = callLine
        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.conv.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.conv.call(call)
                    called = True

            limit -= 1

    def runCmd(self, args):

        call = None
        if args.function:
            call = " ".join(args.function)


        self.run(call, limit = args.limit)

    def clearCmd(self, args):
        self.stack.clear()

    def parseCommand(self, line):


        self.commands.append(line)
        self.update = True
        self.status = ""

        local = self.tryParseLocalVar(line)

        if local:
            line = f"local {line}"
            args = self.parser.parse_args(line.split(" "))
            self.localCmd(local, args)
            self.refreshOutput()

        else:

            try:
                args = self.parser.parse_args(line.split(" "))
                #print(args)

                if args.help:
                    self.showHelp(args.command)
                else:
                    args.func(args)
                    self.refreshOutput()
            except:
                pass
                words = line.split(" ")
                self.showHelp(words[0])


    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...")



    # Initialize the argument parser
    def init_args(self):

        self.parser.add_argument('--help', '-h', action='store_true', help='show this help message and exit')

        subparsers = self.parser.add_subparsers(dest='command', help='sub-command help')

        self.parser.error = nullPrint
        #local command 
        local_parser = subparsers.add_parser('local', help='Declaure a local variable using c syntax')
        local_parser.add_argument('local', nargs='*', help='Local variable declaration')
        local_parser.add_argument('--note','-n', nargs='*', help='Note for the cell')
        local_parser.add_argument('--color', '-c', type=str, help="Color of the cell")
        local_parser.set_defaults(func=self.localCmd)

        # pop command
        pop_parser = subparsers.add_parser('pop', help='pop words from the stack')
        pop_parser.add_argument('count', type=str, nargs='?', help='Number of elements to pop', default=1)
        pop_parser.set_defaults(func=self.popCmd)

        #ret command
        ret_parser = subparsers.add_parser('ret', help='Return (pop, jmp to return address)')
        ret_parser.set_defaults(func=self.retCmd)

        #jmp command
        jmp_parser = subparsers.add_parser('jmp', help='Jump to an address')
        jmp_parser.add_argument('address',  help='Address to jump to')
        jmp_parser.set_defaults(func=self.jmpCmd)
        
        #push command
        push_parser = subparsers.add_parser('push', help='push words to the stack')
        push_parser.add_argument('value', nargs='*', help='Value to push', default='')
        push_parser.add_argument('--size', '-s', type=int, help='Type of value')
        push_parser.add_argument('--label','-l', help='Label for the value')
        push_parser.add_argument('--address','-a', help='Address for the value')
        push_parser.add_argument('--note','-n', nargs='*', help='Note for the value')

        push_parser.set_defaults(func=self.pushCmd)

        #frame command
        frame_parser = subparsers.add_parser('frame', help='Start a new Frame')
        frame_parser.add_argument('name',  help='Name of the frame')
        frame_parser.add_argument('--color', '-c', help='Color of the frame')
        frame_parser.add_argument('--address', '-a', help='address of the frame ')
        frame_parser.add_argument('--size', '-s', type=int, help='size of the frame')
        frame_parser.set_defaults(func=self.frameCmd)

        #call parser
        call_parser = subparsers.add_parser('call', help='Call a function')
        call_parser.add_argument('function', nargs="*", help='Function to call')
        call_parser.set_defaults(func=self.callCmd)

        #function 
        function_parser = subparsers.add_parser('function', help="Define a function using c syntax")
        function_parser.add_argument('function', nargs="*", help='Function signature')
        function_parser.set_defaults(func=self.functionCmd)

        #save 
        save_parser = subparsers.add_parser('save', help='Save the current stack')
        save_parser.add_argument('filename', nargs="?", type=str, help='Filename to save', default='.stackinit')
        save_parser.set_defaults(func=self.saveCmd)

        #load
        load_parser = subparsers.add_parser('load', help='load stack file')
        load_parser.add_argument('filename', type=str, help='Filename to load')
        load_parser.set_defaults(func=self.loadCmd)

        #run 
        run_parser = subparsers.add_parser('run', aliases=['r'], help='Run the a function calls the `call` within a function')
        run_parser.add_argument('function', nargs="*", help='Function call to start, if not provided will use the first call in the current function, or main if no function is active')
        run_parser.add_argument('--limit', '-l', type=int, help='Limit of calls to make', default=20)
        run_parser.set_defaults(func=self.runCmd)

        #write 
        write_parser = subparsers.add_parser('write', aliases=['w'], help='Write data to stack')
        write_parser.add_argument('values', nargs="*", help='Values to write')
        write_parser.add_argument('--address', '-a', type=str, help='Address to write to can be a reference to a cell &<cellName>+4')
        write_parser.add_argument('--count', '-c', type=int, help='Repeated writes', default=1)
        write_parser.add_argument('--file', '-f', type=str, help='binary file to write')
        write_parser.set_defaults(func=self.writeCmd)


        #note
        note_parser = subparsers.add_parser('note', help='Add a note to the stack')
        note_parser.add_argument('note', nargs="*", help='Note to add')
        note_parser.add_argument('--color', '-c', type=str, help='Color of the note', default='red')
        note_parser.add_argument('--address', '-a', type=str, help='Address of the note')
        note_parser.set_defaults(func=self.noteCmd)

        #leave 
        leave_parser = subparsers.add_parser('leave', help='Pops words from stack until the previous base pointer')
        leave_parser.set_defaults(func=self.leaveCmd)


        clear_parser = subparsers.add_parser('clear', help='Clears the stack')
        clear_parser.set_defaults(func=self.clearCmd)

        for sub in subparsers.choices:
            subparsers.choices[sub].print_help =  nullPrint
            subparsers.choices[sub].error = nullPrint


        self.subparsers = subparsers
