import argparse
import subprocess
import sys
import tempfile
import os.path
from pathlib import Path

import colorama
from functools import wraps
# once Python3.10 is mainstream, we can replace Union[str, None] by str | None
from typing import Union, Callable, Any
from prompt_toolkit.completion import Completer, NestedCompleter, WordCompleter
from prompt_toolkit.completion.base import Completion
from prompt_toolkit.document import Document

from chameleon_enum import Status

# Colorama shorthands
CR = colorama.Fore.RED
CG = colorama.Fore.GREEN
CB = colorama.Fore.BLUE
CC = colorama.Fore.CYAN
CY = colorama.Fore.YELLOW
CM = colorama.Fore.MAGENTA
C0 = colorama.Style.RESET_ALL


def get_resource_dir(relative_path: str):
    """
    Get the resource directory of the program.
    Returns the temporary directory where files are extracted after being packaged with PyInstaller, or the directory where the script is located in the development environment.
    """
    if getattr(sys, 'frozen', False) and hasattr(sys, '_MEIPASS'):
        base_dir = Path(sys._MEIPASS)
    else:
        base_dir = Path(__file__).parent
    return base_dir / relative_path


default_cwd = get_resource_dir("bin")


class ArgsParserError(Exception):
    pass


class ParserExitIntercept(Exception):
    pass


class UnexpectedResponseError(Exception):
    """
    Unexpected response exception
    """


class ArgumentParserNoExit(argparse.ArgumentParser):
    """
        If arg ArgumentParser parse error, we can't exit process,
        we must raise exception to stop parse
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.add_help = False
        self.description = "Please enter correct parameters"
        self.help_requested = False

    def exit(self, status: int = 0, message: Union[str, None] = None):
        if message:
            raise ParserExitIntercept(message)
            # status=0 means help was printed; raise to stop argparse continuing
            # to validate required args (which would cause a second print_help call)
        raise ParserExitIntercept('')

    def error(self, message: str):
        args = {'prog': self.prog, 'message': message}
        raise ArgsParserError('%(prog)s: error: %(message)s\n' % args)


def print_help(self):
    """
    Colorize argparse help
    """
    print("-" * 80)
    print(color_string((CR, self.prog)))

    # Get the help text and split it, filtering out leading empty lines
    raw_lines = self.format_help().splitlines()
    lines = [line for line in raw_lines if line.strip() or line == '']

    # Find the usage block safely
    usage_start = -1
    for i, line in enumerate(lines):
        if line.strip().startswith('usage:'):
            usage_start = i
            break

    if usage_start != -1:
        # We found a usage line, extract the block until the first empty line
        try:
            empty_after_usage = lines.index('', usage_start)
            usage = lines[usage_start:empty_after_usage]

            # Apply coloring to the usage string
            usage[0] = usage[0].replace('usage:', f'{color_string((CG, "usage:"))}\n ')
            usage[0] = usage[0].replace(self.prog, color_string((CR, self.prog)))

            # Reformat indentation and print
            usage_to_print = [usage[0]] + [x[4:] for x in usage[1:]] + ['']
            print('\n'.join(usage_to_print))

            # Advance lines pointer to after the usage block
            lines = lines[empty_after_usage + 1:]
        except ValueError:
            # If no empty line found, just print what we have
            print('\n'.join(lines[usage_start:]))
            lines = []

    # Print description if available
    if lines and lines[0].strip() != '':
        try:
            desc_end = lines.index('')
            desc = lines[:desc_end]
            print(color_string((CC, "\n".join(desc))))
            lines = lines[desc_end + 1:]
        except ValueError:
            pass

    # Handle options and positional arguments without crashing on strict matches
    for line in lines:
        clean_line = line.strip().lower()
        if clean_line == 'positional arguments:':
            print(color_string((CG, line)))
        elif clean_line in ['options:', 'optional arguments:']:
            print(color_string((CG, line)))
        else:
            print(line)

    print('')
    self.help_requested = True


def print_mem_dump(bindata, blocksize):

    hexadecimal_len = blocksize*3+1
    ascii_len = blocksize+1
    print(f"[=] ----+{hexadecimal_len*'-'}+{ascii_len*'-'}")
    print(f"[=] blk | data{(hexadecimal_len-5)*' '}| ascii")
    print(f"[=] ----+{hexadecimal_len*'-'}+{ascii_len*'-'}")

    blocks = [bindata[i:i+blocksize] for i in range(0, len(bindata), blocksize)]
    blk_index = 1
    for b in blocks:
        hexstr = ' '.join(b.hex()[i:i+2] for i in range(0, len(b.hex()), 2))
        asciistr = ''.join([chr(b[i]) if (b[i] > 31 and b[i] < 127) else '.' for i in range(0, len(b), 1)])
        print(f"[=] {blk_index:3} | {hexstr.upper()} | {asciistr} ")
        blk_index += 1


def print_key_table(key_map):
    key_width = max(
        max(len(k) for k in key_map["A"].values()),
        max(len(k) for k in key_map["B"].values()),
        len("key A"),
        len("key B"),
    )
    header_line = f"[=] {'-'*5}+{'-'*(key_width+2)}+{'-'*(key_width+2)}"
    print(header_line)
    print(f"[=]  sec | key A{' '*(key_width-5)} | key B{' '*(key_width-5)}")
    print(header_line)
    for sec, (a, b) in enumerate(zip(key_map["A"].values(), key_map["B"].values())):
        print(f"[=]  {sec:02d}  | {a:{key_width}} | {b:{key_width}}")
    print(header_line)


def _swap_endian(x):
    x = ((x >> 8) & 0x00ff00ff) | ((x & 0x00ff00ff) << 8)
    x = (x >> 16) | (x << 16)
    return x & 0xFFFFFFFF


def prng_successor(x, n):
    x = _swap_endian(x)

    while n > 0:
        x = (x >> 1) | (
            (((x >> 16) ^ (x >> 18) ^ (x >> 19) ^ (x >> 21)) & 0x1) << 31
        )
        x = x & 0xFFFFFFFF
        n -= 1

    return _swap_endian(x)


def reconstruct_full_nt(response_data, offset):
    nt = int.from_bytes(response_data[offset: offset + 2], byteorder='big')

    return (nt << 16) | prng_successor(nt, 16)


def parity_to_str(nt_par_err):
    return "".join(
        [
            str((nt_par_err >> 3) & 1),
            str((nt_par_err >> 2) & 1),
            str((nt_par_err >> 1) & 1),
            str(nt_par_err & 1),
        ]
    )

def odd_parity_byte(value: int) -> int:
    assert(value >=0 and value <= 255)
    for i in [4, 2, 1]:
        value ^= value >> i
    return (value & 1) ^ 1

def execute_tool(tool_name, args):
    if sys.platform == "win32":
        tool_executable = f"{tool_name}.exe"
    else:
        tool_executable = f"./{tool_name}"

    tool_path = os.path.join(default_cwd, tool_executable)
    cmd_recover_list = [tool_path]
    cmd_recover_list.extend(args)

    # print(f"Executing: {' '.join(cmd_recover_list)}")

    temp_output_file = tempfile.NamedTemporaryFile(
        suffix=".log", prefix="output_", delete=True,
        mode='w+', encoding='utf-8', errors='replace'
    )

    process = subprocess.Popen(
        cmd_recover_list,
        cwd=tempfile.gettempdir(),
        stdout=temp_output_file,
        stderr=subprocess.STDOUT,
    )

    ret_code = process.wait()
    temp_output_file.seek(0)

    if ret_code:
        raise Exception('Failed to execute tool: ' + temp_output_file.read())

    return temp_output_file.read()


def tqdm_if_exists(iterator):
    try:
        import tqdm
        return tqdm.tqdm(iterator)
    except ImportError:
        return iterator


def expect_response(accepted_responses: Union[int, list[int]]) -> Callable[..., Any]:
    """
    Decorator for wrapping a Chameleon CMD function to check its response
    for expected return codes and throwing an exception otherwise
    """
    if isinstance(accepted_responses, int):
        accepted_responses = [accepted_responses]

    def decorator(func):
        @wraps(func)
        def error_throwing_func(*args, **kwargs):
            ret = func(*args, **kwargs)
            if ret.status not in accepted_responses:
                try:
                    status_string = str(Status(ret.status))
                except ValueError:
                    status_string = f"Unexpected response and unknown status {ret.status}"
                raise UnexpectedResponseError(status_string)

            return ret.parsed

        return error_throwing_func

    return decorator


def color_string(*args):
    result = []
    for arg in args:
        result.append(f"{arg[0]}{arg[1]}")
    result.append(C0)
    return "".join(result)


class CLITree:
    """
    Class holding a

    :param name: Name of the command (e.g. "set")
    :param help_text: Hint displayed for the command
    :param fullname: Full name of the command that includes previous commands (e.g. "hw settings animation")
    :param cls: A BaseCLIUnit instance handling the command
    """

    def __init__(self, name: str = "", help_text: Union[str, None] = None, fullname: Union[str, None] = None,
                 children: Union[list["CLITree"], None] = None, cls=None, root=False) -> None:
        self.name = name
        self.help_text = help_text
        self.fullname = fullname if fullname else name
        self.children = children if children else list()
        self.cls = cls
        self.root = root
        if self.help_text is None and not root:
            assert self.cls is not None
            parser = self.cls().args_parser()
            assert parser is not None
            self.help_text = parser.description

    def subgroup(self, name, help_text=None):
        """
        Create a child command group

        :param name: Name of the command group
        :param help_text: Hint displayed for the group
        """
        child = CLITree(
            name=name,
            fullname=f'{self.fullname} {name}' if not self.root else f'{name}',
            help_text=help_text)
        self.children.append(child)
        return child

    def command(self, name):
        """
        Create a child command

        :param name: Name of the command
        """
        def decorator(cls):
            self.children.append(CLITree(
                name=name,
                fullname=f'{self.fullname} {name}' if not self.root else f'{name}',
                cls=cls))
            return cls
        return decorator


class CustomNestedCompleter(NestedCompleter):
    """
    Copy of the NestedCompleter class that accepts a CLITree object and
    supports meta_dict for descriptions
    """

    def __init__(
        self, options, ignore_case: bool = True, meta_dict: dict = {}
    ) -> None:
        self.options = options
        self.ignore_case = ignore_case
        self.meta_dict = meta_dict

    def __repr__(self) -> str:
        return f"CustomNestedCompleter({self.options!r}, ignore_case={self.ignore_case!r})"

    @classmethod
    def from_clitree(cls, node):
        options = {}
        meta_dict = {}

        for child_node in node.children:
            if child_node.cls:
                # CLITree is a standalone command with arguments
                options[child_node.name] = ArgparseCompleter(
                    child_node.cls().args_parser())
            else:
                # CLITree is a command group
                options[child_node.name] = cls.from_clitree(child_node)
                meta_dict[child_node.name] = child_node.help_text

        return cls(options, meta_dict=meta_dict)

    def get_completions(self, document, complete_event):
        # Split document.
        text = document.text_before_cursor.lstrip()
        stripped_len = len(document.text_before_cursor) - len(text)

        # If there is a space, check for the first term, and use a sub_completer.
        if " " in text:
            first_term = text.split()[0]
            completer = self.options.get(first_term)

            # If we have a sub completer, use this for the completions.
            if completer is not None:
                remaining_text = text[len(first_term):].lstrip()
                move_cursor = len(text) - len(remaining_text) + stripped_len

                new_document = Document(
                    remaining_text,
                    cursor_position=document.cursor_position - move_cursor,
                )

                yield from completer.get_completions(new_document, complete_event)

        # No space in the input: behave exactly like `WordCompleter`.
        else:
            completer = WordCompleter(
                list(self.options.keys()), ignore_case=self.ignore_case, meta_dict=self.meta_dict
            )
            yield from completer.get_completions(document, complete_event)


class ArgparseCompleter(Completer):
    """
    Completer instance for autocompletion of ArgumentParser arguments

    :param parser: ArgumentParser instance
    """

    def __init__(self, parser) -> None:
        self.parser: ArgumentParserNoExit = parser

    def check_tokens(self, parsed, unparsed):
        suggestions = {}

        def check_arg(tokens):
            return tokens and tokens[0].startswith('-')

        if not parsed and not unparsed:
            # No tokens detected, just show all flags
            for action in self.parser._actions:
                for opt in action.option_strings:
                    suggestions[opt] = action.help
            return [], [], suggestions

        token = unparsed.pop(0)

        for action in self.parser._actions:
            if any(opt == token for opt in action.option_strings):
                # Argument fully matches the token
                parsed.append(token)

                if action.choices:
                    # Autocomplete with choices
                    if unparsed:
                        # Autocomplete values
                        value = unparsed.pop(0)
                        for choice in action.choices:
                            if str(choice).startswith(value):
                                suggestions[str(choice)] = None

                        parsed.append(value)

                        if check_arg(unparsed):
                            parsed, unparsed, suggestions = self.check_tokens(
                                parsed, unparsed)

                    else:
                        # Show all possible values
                        for choice in action.choices:
                            suggestions[str(choice)] = None

                    break
                else:
                    # No choices, process further arguments
                    if check_arg(unparsed):
                        parsed, unparsed, suggestions = self.check_tokens(
                            parsed, unparsed)
                    break
            elif any(opt.startswith(token) for opt in action.option_strings):
                for opt in action.option_strings:
                    if opt.startswith(token):
                        suggestions[opt] = action.help

        if suggestions:
            unparsed.insert(0, token)

        return parsed, unparsed, suggestions

    def get_completions(self, document, complete_event):
        text = document.text_before_cursor
        word_before_cursor = document.text_before_cursor.split(' ')[-1]

        _, _, suggestions = self.check_tokens(list(), text.split())

        for key, suggestion in suggestions.items():
            yield Completion(key, -len(word_before_cursor), display=key, display_meta=suggestion)