#!/bin/env python3
from prompt_toolkit.application import Application
from prompt_toolkit.key_binding import KeyBindings
from prompt_toolkit.completion import WordCompleter, PathCompleter, Completer
from prompt_toolkit.layout.containers import HSplit, VSplit, Window, ScrollOffsets, FloatContainer,Float
from prompt_toolkit.layout.layout import Layout
from prompt_toolkit.layout import ScrollablePane
from prompt_toolkit.layout.controls import FormattedTextControl, BufferControl
from prompt_toolkit.mouse_events import MouseButton, MouseEvent, MouseEventType
from prompt_toolkit.widgets import TextArea, Frame, Label
from prompt_toolkit.history import FileHistory
from prompt_toolkit.buffer import Buffer
from prompt_toolkit.widgets import HorizontalLine
from prompt_toolkit.data_structures import Point
from prompt_toolkit.document import Document
from prompt_toolkit.layout.menus import CompletionsMenu
from prompt_toolkit import print_formatted_text, ANSI
from prompt_toolkit.completion import NestedCompleter


import asyncio

import os
import argparse
import typing
import importlib.metadata
import builtins


from contextlib import contextmanager


from staq.stackSession import StackSession


args = None
parser = None
session = StackSession()

version = importlib.metadata.version("staq")

completer = NestedCompleter.from_nested_dict(
    {
        "pop": None,
        "push" : {"--size": None, "--label": None, "--note": None}
    }
)

completion_lists = {
    "pop" : [],
    "ret" : [],
    "push" : ["--size ", "--label ", "--note ", "--address "],
    "frame" : ['--color '],
    "call": [],
    "function" : [],
    "load": [],
    "write": ["--address ", "--file "]
}



class CommandCompleter(Completer):
    def __init__(self):
        self.path_completer = PathCompleter()
        self.args_completer = WordCompleter( [],ignore_case=True)
        self.base_completer = WordCompleter( list(completion_lists.keys()),ignore_case=True)

    def get_completions(self, document: Document, complete_event):
        text = document.text

        commands = text.split(";")

        if len(commands) == 0:
            yield completion_lists.keys()
    
        words = commands[-1].split()

        if len(words) == 0:
            yield from self.base_completer.get_completions(document, complete_event)
        else:

            cmd = words[0]
            last = words[-1]

            currentArg = ""
            midWord = True

            if commands[-1].endswith(" "):
                midWord = False

            for word in words:
                if word.startswith("--"):
                    currentArg = word


            if len(words) > 1 and (words[0] == 'load' and (len(words) <= 2) and midWord) or currentArg == '--file':
                # Apply path completer to the last word
                path_text = document.get_word_before_cursor(WORD=True)
                path_document = Document(path_text, len(path_text))
                yield from self.path_completer.get_completions(path_document, complete_event)
            elif cmd == 'call' or cmd == 'run':

                
                self.args_completer.words = session.functions.keys()
                yield from self.args_completer.get_completions(document, complete_event)
            elif cmd in completion_lists:
                self.args_completer.words = completion_lists[cmd]
                yield from self.args_completer.get_completions(document, complete_event)
            else:
                # Default completions for other cases (if any)
                yield from self.base_completer.get_completions(document, complete_event)
        
command_completer = CommandCompleter()

banner_text =f"Staq v{version}"



def accept(buff):
    global command_completer

    line = buff.text

    commands = line.split(';')

    for cmd in commands:
            
        if cmd != "":
            session.parseCommand(cmd)
            if session.update:
                refreshOutput()
            top_info.text = session.status

    if line == "":
        refreshOutput()

# Define key bindings
kb = KeyBindings()

# Exit application
@kb.add('c-c')
def exit_app(event):
    
    event.app.exit(0)

@kb.add('c-s')
def save_stack(event):
    session.parseCommand("save")
    pass




def print_info(event):
    util_info.text = str(input_field.completer.words)
    pass


# Banner content

banner_left = Label(banner_text, style="bg:ansiblack fg:ansigreen")

top_info = Label("", style="bg:ansiblack fg:ansired")
util_info = Label("", style="bg:ansiblack fg:ansiblue")

banner_right = HSplit([top_info])

#banner_right = Label("SimTime: 000000\nrepli-4h23\nconnected to .sdf.sdf", style="bg:ansiblack fg:ansired")

banner_content = VSplit([banner_left, banner_right])
# Divider line
divider = HorizontalLine()
divider.window.style = "fg:ansigreen"

# Main content area - initially empty, will be updated dynamically
_line_count =0
outputField = TextArea(multiline=True, style="bg:ansiblack fg:ansigreen")

def getCursPos():
    global _line_count
    return Point(x=0, y=_line_count)


output_control = FormattedTextControl( text=[('ansigreen', '')], get_cursor_position=getCursPos, focusable=True)
output_window = Window( content=output_control, style="class:main-content", wrap_lines= False, allow_scroll_beyond_bottom= True, scroll_offsets=ScrollOffsets(top=10000, bottom=10000))


# Input area
home_path = os.path.expanduser('~')
cmdHistory = FileHistory(home_path+'/.staq-history')

input_field = Buffer( multiline=False, complete_while_typing=True,history=cmdHistory, completer=command_completer)
input_window = Window(BufferControl(buffer=input_field), height=2, style='class:input-win')
prompt_window = Window(FormattedTextControl(text="$>"), height=1, style='fg:ansigreen', width=3)

input_field.accept_handler = accept


reg_control = FormattedTextControl( text=[('ansigreen', '')], get_cursor_position=getCursPos, focusable=True)

floating_window = Window(content=reg_control, style="class:main-content", wrap_lines= False, height=5, width=25) 

# Layout definition
layout = FloatContainer(
       content= HSplit([
        banner_content,  # Banner at the top
        divider,         # Divider line
        output_window,    # Main content area
        divider,         # Divider line
        # input_window,     # Input area at the bottom
        VSplit([prompt_window, input_window]) # Input area at the bottom
    ]),
    floats=[
        Float(xcursor=True,
              ycursor=True,
              content=CompletionsMenu(max_height=16, scroll_offset=1)),
        Float(top=2,
              right=1,
              content=floating_window)
    ]
)
# layout = HSplit([
#             banner_content,  # Banner at the top
#             divider,         # Divider line
#             output_window,    # Main content area
#             divider,         # Divider line
#             # input_window,     # Input area at the bottom
#             VSplit([prompt_window, input_window]) # Input area at the bottom
#         ])



# Initialize the argument parser
def init_args():
    global parser
    parser = argparse.ArgumentParser("Tool to interact with replicant message broker")
    parser.add_argument('-b', '--base-address', type=int, help='Stack Base Address', default=0xffff)
    parser.add_argument('-f', '--file', type=str, help='Stack File to load', default="")
    parser.add_argument('-i', '--init', type=str, help='Initial Stack File to load', default=".stack.lst")
    parser.add_argument('-x', '--execute', nargs="*", type=str, help='Execute command', default="")




# Application
application = Application(layout=Layout(layout, focused_element=input_field), key_bindings=kb, full_screen=True ,mouse_support=True)

def refreshOutput():
    global prevRegLines
    
    text = session.stack.toAnsi(showAddress=True)
    regText = session.conv.registersToAnsi(width=floating_window.width -2)
    regAnsi = ANSI(regText)
    ansi = ANSI(text)

    output_control.text = ansi
    reg_control.text = regAnsi






def printToConsole( text : typing.Union[str, list, tuple], color = 'ansigreen', mewline = True):
    
    
    global _line_count



    current_content = output_control.text

    if type(text) is str: 
       
       if(mewline): #Plain text
            text += "\n"
       _line_count += text.count("\n")
       output_control.text += [(color, text )]

    elif type(text) is list: #List of color,string tuples
        for tup in text:
            _line_count += tup[1].count("\n")
        
        if(mewline):
            text.append(('',"\n"))
            _line_count+=1

        output_control.text += text
    elif type(text) is tuple:   #Single color,string tuple
        _line_count += text[1].count("\n")

        if(mewline):
            text = [text, ('',"\n")]
            _line_count+=1
        output_control.text += text

    
    output_window.vertical_scroll = _line_count
    application.invalidate()
       
       
    pass

outputText = ""

def printRedirect(*args, **kwargs):
    global outputText
    output = " ".join(map(str, args))

    outputText+= output

    ansi = ANSI(outputText)
    output_control.text = ansi


def clearConsole():
    global _line_count
    _line_count = 0
    output_control.text = [('ansigreen', '')]
    application.invalidate()

    pass


async def run():
    global parser
    global args
    global input_field
    global command_completer
    global session 

    init_args()
    args= parser.parse_args()

    #session.loadYaml("test/test1.stack.yml")

    session.stack.setBaseAddress( args.base_address)
    

    #get list of subparsers from session parser 

    builtins.print = printRedirect
    session.print = printRedirect

    if args.file != "":
        session.loadYaml(args.file)

    if os.path.exists(args.init):
        with open(args.init, 'r') as f:
            for line in f:
                session.parseCommand(line)

    if args.execute != "":
        line  = " ".join(args.execute)
        cmds = line.split(";")
        for cmd in cmds:
            session.parseCommand(cmd)

    


    
    refreshOutput()



    result = await application.run_async()



if __name__ == '__main__':
   asyncio.run( run())