import t, c
import string
import memhub
import stdio
import hashtable


# ============================================================
# argparse.vp — Command line argument parsing library, inspired by Python argparse
#
# Usage:
#   argparse._mbuddy = mb
#   parser = argparse.ArgumentParser("myprog", "A test", pool=mb)
#   parser.add_argument("--input", "-i", arg_type=argparse.STRING,
#                       required=True, help="input file")
#   parser.add_argument("--count", "-c", arg_type=argparse.INT,
#                       default=0, help="repeat count")
#   parser.add_argument("--verbose", "-v", action=argparse.STORE_TRUE,
#                       help="verbose mode")
#   parser.add_argument("output", arg_type=argparse.STRING,
#                       help="output path")
#   args = parser.parse_args(argc, argv)
#   if args is not None:
#       inp: str = args.get_str("input")
#       cnt: int = args.get_int("count")
#       verb: bool = args.get_bool("verbose")
#       out: str = args.get_str("output")
#       argparse.release(args)
# ============================================================

# Type constants
STRING: t.CDefine = 0
INT: t.CDefine = 1
BOOL: t.CDefine = 2

# Action constants
STORE: t.CDefine = 0
STORE_TRUE: t.CDefine = 1
STORE_FALSE: t.CDefine = 2
COUNT: t.CDefine = 3

# Limits
MAX_ARGS: t.CDefine = 32

# Module-level membuddy pointer (assigned inside main by user)
_mbuddy: memhub.MemBuddy | t.CPtr


# ============================================================
# Argument — Argument definition structure
# ============================================================

class Argument:
    name: str               # "--input" or positional "output"
    short_name: str         # "-i" or null
    help_text: str          # Help description text
    arg_type: int           # STRING, INT, BOOL
    action: int             # STORE, STORE_TRUE, STORE_FALSE, COUNT
    default_int: int        # Default value for integer types
    default_str: str        # Default value for string types
    required: bool          # Whether this argument is mandatory
    is_positional: bool     # True for positional parameters
    dest: str               # Lookup key ("--input" → "input")


# ============================================================
# ParsedArgs — Parsing result container
# ============================================================

class ParsedArgs:
    _ht: hashtable.HashTable | t.CPtr
    __mbuddy__: memhub.MemBuddy | t.CPtr

    def __new__(self, mb: memhub.MemBuddy | t.CPtr) -> t.CPtr:
        return mb.alloc(16)

    def __init__(self, mb: memhub.MemBuddy | t.CPtr):
        self.__mbuddy__ = mb
        self._ht = hashtable.HashTable(mb)

    def get_str(self, name: str) -> str:
        """Fetch string value. Returns None if key not found."""
        if self._ht.__contains__(name) == 0:
            return None
        return self._ht[name]

    def get_int(self, name: str) -> int:
        """Fetch integer value. Returns 0 if key not found."""
        if self._ht.__contains__(name) == 0:
            return 0
        ptr: t.CPtr = self._ht[name]
        return c.Deref(t.CInt(t.CUInt64T(ptr), t.CPtr))

    def get_bool(self, name: str) -> bool:
        """Fetch boolean value. Returns false if key not found."""
        if self._ht.__contains__(name) == 0:
            return False
        ptr: t.CPtr = self._ht[name]
        return c.Deref(t.CInt(t.CUInt64T(ptr), t.CPtr)) != 0

    def has(self, name: str) -> bool:
        """Check whether specified key exists."""
        return self._ht.__contains__(name) != 0

    def _set_str(self, name: str, val: str):
        self._ht.set_str(name, val)

    def _set_int(self, name: str, val: int):
        self._ht.set_int(name, val)

    def _set_bool(self, name: str, val: bool):
        if val:
            self._ht.set_int(name, 1)
        else:
            self._ht.set_int(name, 0)


# ============================================================
# ArgumentParser — Command line argument parser
# ============================================================

class ArgumentParser:
    _prog: str
    _description: str
    _args: Argument | t.CPtr
    _arg_count: int
    __mbuddy__: memhub.MemBuddy | t.CPtr

    def __new__(self, prog: str = "", description: str = "",
                pool: memhub.MemBuddy | t.CPtr = None) -> t.CPtr:
        stdio.printf("[DBG-argparse] __new__ enter pool=%p\n", pool)
        mb: memhub.MemBuddy | t.CPtr = pool
        if mb is None:
            mb = _mbuddy
        stdio.printf("[DBG-argparse] __new__ mb=%p\n", mb)
        ret: t.CPtr = mb.alloc(40)
        stdio.printf("[DBG-argparse] __new__ ret=%p\n", ret)
        return ret

    def __init__(self, prog: str = "", description: str = "",
                 pool: memhub.MemBuddy | t.CPtr = None):
        stdio.printf("[DBG-argparse] __init__ enter self=%p pool=%p\n", self, pool)
        mb: memhub.MemBuddy | t.CPtr = pool
        if mb is None:
            mb = _mbuddy
        stdio.printf("[DBG-argparse] __init__ mb=%p\n", mb)
        self.__mbuddy__ = mb
        self._prog = prog
        self._description = description
        self._arg_count = 0
        stdio.printf("[DBG-argparse] __init__ before alloc\n")
        self._args = mb.alloc(MAX_ARGS * Argument.__sizeof__())
        stdio.printf("[DBG-argparse] __init__ after alloc _args=%p\n", self._args)

    def add_argument(self, name: str = "", short: str = None,
                     arg_type: int = 0, default: int = 0,
                     default_str: str = None, required: bool = False,
                     action: int = 0, help: str = None):
        """Register a command line argument.

        Args:
            name: Argument name ("--input" or positional "output")
            short: Short option ("-i"), optional
            arg_type: Data type (STRING/INT/BOOL)
            default: Default integer value
            default_str: Default string value
            required: Mandatory flag
            action: Action mode (STORE/STORE_TRUE/STORE_FALSE/COUNT)
            help: Help description text
        """
        stdio.printf("[DBG-argparse] enter add_argument name=%s args_ptr=%p\n", name, self._args)
        if self._arg_count >= MAX_ARGS:
            return
        idx: int = self._arg_count
        stdio.printf("[DBG-argparse] idx=%d sizeof(Argument)=%d\n", idx, Argument.__sizeof__())
        arg: Argument | t.CPtr = t.CPtr(t.CUInt64T(self._args) + idx * Argument.__sizeof__())
        stdio.printf("[DBG-argparse] arg ptr=%p\n", arg)
        arg.name = name
        stdio.printf("[DBG-argparse] after arg.name=%s\n", arg.name)
        arg.short_name = short
        arg.help_text = help
        arg.arg_type = arg_type
        arg.action = action
        arg.default_int = default
        arg.default_str = default_str
        arg.required = required
        stdio.printf("[DBG-argparse] before dest logic\n")
        # Determine positional flag
        if name is not None and name[0] != '-':
            arg.is_positional = True
            arg.dest = name
        else:
            arg.is_positional = False
            if name is not None and name[0] == '-' and name[1] == '-':
                arg.dest = name + 2
            else:
                arg.dest = name + 1
        stdio.printf("[DBG-argparse] after dest logic\n")
        self._arg_count = idx + 1
        stdio.printf("[DBG-argparse] exit add_argument\n")

    def _get_arg(self, idx: int) -> Argument | t.CPtr:
        return t.CPtr(t.CUInt64T(self._args) + idx * Argument.__sizeof__())

    def _find_long(self, token: str) -> int:
        """Exact match for long option name"""
        i: int = 0
        while i < self._arg_count:
            arg: Argument | t.CPtr = self._get_arg(i)
            if string.strcmp(arg.name, token) == 0:
                return i
            i += 1
        return -1

    def _find_long_prefix(self, token: str) -> int:
        """Match prefix for --name=value syntax"""
        i: int = 0
        while i < self._arg_count:
            arg: Argument | t.CPtr = self._get_arg(i)
            name_len: int = string.strlen(arg.name)
            if string.strncmp(token, arg.name, name_len) == 0:
                if token[name_len] == '=':
                    return i
            i += 1
        return -1

    def _find_short(self, token: str) -> int:
        i: int = 0
        while i < self._arg_count:
            arg: Argument | t.CPtr = self._get_arg(i)
            if arg.short_name is not None:
                if string.strcmp(arg.short_name, token) == 0:
                    return i
            i += 1
        return -1

    def _find_positional(self, pos_idx: int) -> int:
        count: int = 0
        i: int = 0
        while i < self._arg_count:
            arg: Argument | t.CPtr = self._get_arg(i)
            if arg.is_positional:
                if count == pos_idx:
                    return i
                count += 1
            i += 1
        return -1

    def _store_value(self, args: ParsedArgs | t.CPtr,
                     arg: Argument | t.CPtr, value: str):
        if arg.arg_type == INT:
            args._set_int(arg.dest, string.atoi(value))
        elif arg.arg_type == STRING:
            args._set_str(arg.dest, value)
        elif arg.arg_type == BOOL:
            v: int = string.atoi(value)
            args._set_bool(arg.dest, v != 0)

    def _apply_action(self, args: ParsedArgs | t.CPtr, arg: Argument | t.CPtr):
        if arg.action == STORE_TRUE:
            args._set_bool(arg.dest, True)
        elif arg.action == STORE_FALSE:
            args._set_bool(arg.dest, False)
        elif arg.action == COUNT:
            cur: int = args.get_int(arg.dest)
            args._set_int(arg.dest, cur + 1)

    def parse_args(self, argc: int, argv: str | t.CPtr) -> ParsedArgs | t.CPtr:
        """Parse command line arguments.

        Args:
            argc: Argument count
            argv: Argument array (char**, argv[0] = program name)
        Returns:
            Pointer to ParsedArgs; returns None on failure
        """
        if self.__mbuddy__ is None:
            return None

        args: ParsedArgs | t.CPtr = ParsedArgs(self.__mbuddy__)
        if args is None:
            return None

        # Initialize default values
        i: int = 0
        while i < self._arg_count:
            arg: Argument | t.CPtr = self._get_arg(i)
            if arg.action == STORE:
                if arg.arg_type == INT:
                    args._set_int(arg.dest, arg.default_int)
                elif arg.arg_type == STRING:
                    args._set_str(arg.dest, arg.default_str)
            elif arg.action == STORE_TRUE:
                args._set_bool(arg.dest, False)
            elif arg.action == STORE_FALSE:
                args._set_bool(arg.dest, True)
            elif arg.action == COUNT:
                args._set_int(arg.dest, 0)
            i += 1

        # Traverse argv
        idx: int = 1
        pos_idx: int = 0
        token: str = None
        found: int = 0
        arg: Argument | t.CPtr = None
        value: str = None
        entry: str | t.CPtr = None

        while idx < argc:
            entry = t.CPtr(t.CUInt64T(argv) + idx * 8)
            token = entry[0]
            if token is None:
                break

            # "--" terminator: all subsequent tokens become positional arguments
            if token[0] == '-' and token[1] == '-' and token[2] == 0:
                idx += 1
                while idx < argc:
                    entry = t.CPtr(t.CUInt64T(argv) + idx * 8)
                    token = entry[0]
                    found = self._find_positional(pos_idx)
                    if found < 0:
                        stdio.printf("error: unexpected argument '%s'\n", token)
                        release(args)
                        return None
                    arg = self._get_arg(found)
                    self._store_value(args, arg, token)
                    pos_idx += 1
                    idx += 1
                break

            # Long option --name or --name=value
            if token[0] == '-' and token[1] == '-':
                has_eq: bool = False
                p: str = token
                while p[0] != 0:
                    if p[0] == '=':
                        has_eq = True
                        break
                    p += 1

                if has_eq:
                    # --name=value syntax
                    found = self._find_long_prefix(token)
                    if found < 0:
                        stdio.printf("error: unrecognized argument '%s'\n", token)
                        release(args)
                        return None
                    arg = self._get_arg(found)
                    arg_name_len: int = string.strlen(arg.name)
                    value = token + arg_name_len + 1
                    self._store_value(args, arg, value)
                    idx += 1
                else:
                    # --name with value in next token or flag action
                    found = self._find_long(token)
                    if found < 0:
                        stdio.printf("error: unrecognized argument '%s'\n", token)
                        release(args)
                        return None
                    arg = self._get_arg(found)
                    if arg.action == STORE:
                        if idx + 1 >= argc:
                            stdio.printf("error: %s requires a value\n", arg.name)
                            release(args)
                            return None
                        entry = t.CPtr(t.CUInt64T(argv) + (idx + 1) * 8)
                        value = entry[0]
                        self._store_value(args, arg, value)
                        idx += 2
                    else:
                        self._apply_action(args, arg)
                        idx += 1

            # Short option -x
            elif token[0] == '-':
                found = self._find_short(token)
                if found < 0:
                    stdio.printf("error: unrecognized argument '%s'\n", token)
                    release(args)
                    return None
                arg = self._get_arg(found)
                if arg.action == STORE:
                    if idx + 1 >= argc:
                        stdio.printf("error: %s requires a value\n", arg.name)
                        release(args)
                        return None
                    entry = t.CPtr(t.CUInt64T(argv) + (idx + 1) * 8)
                    value = entry[0]
                    self._store_value(args, arg, value)
                    idx += 2
                else:
                    self._apply_action(args, arg)
                    idx += 1

            # Positional argument
            else:
                found = self._find_positional(pos_idx)
                if found < 0:
                    stdio.printf("error: unexpected argument '%s'\n", token)
                    release(args)
                    return None
                arg = self._get_arg(found)
                self._store_value(args, arg, token)
                pos_idx += 1
                idx += 1

        # Validate required arguments
        i = 0
        while i < self._arg_count:
            arg = self._get_arg(i)
            if arg.required:
                if not args.has(arg.dest):
                    stdio.printf("error: missing required argument '%s'\n", arg.name)
                    release(args)
                    return None
            i += 1

        return args

    def print_help(self):
        """Print complete help message"""
        stdio.printf("usage: %s [options]\n", self._prog)
        if self._description is not None and self._description[0] != 0:
            stdio.printf("\n%s\n", self._description)
        stdio.printf("\npositional arguments:\n")
        i: int = 0
        while i < self._arg_count:
            arg: Argument | t.CPtr = self._get_arg(i)
            if arg.is_positional:
                ht: str = arg.help_text
                if ht is None:
                    ht = ""
                stdio.printf("  %-20s %s\n", arg.name, ht)
            i += 1
        stdio.printf("\noptional arguments:\n")
        i = 0
        while i < self._arg_count:
            arg = self._get_arg(i)
            if not arg.is_positional:
                ht: str = arg.help_text
                if ht is None:
                    ht = ""
                if arg.short_name is not None:
                    stdio.printf("  %s, %-18s %s\n", arg.short_name, arg.name, ht)
                else:
                    stdio.printf("  %-22s %s\n", arg.name, ht)
            i += 1

    def print_usage(self):
        """Print brief usage summary"""
        stdio.printf("usage: %s [options]\n", self._prog)


# ============================================================
# release — Release ParsedArgs and all internal resources
# ============================================================

def release(args: ParsedArgs | t.CPtr):
    """Free ParsedArgs object and its inner HashTable resources."""
    if args is None:
        return
    if args.__mbuddy__ is None:
        return
    if args._ht is not None:
        # 用临时变量拆分链式属性访问，避免 args._ht.__slots__ 跨模块类型推断失败
        # 注意：类型注解必须用 hashtable.HashTable（带模块前缀），
        # 否则编译器无法解析 HashTable 类型，导致 ht.__slots__ 属性访问失败
        ht: hashtable.HashTable | t.CPtr = args._ht
        if ht.__slots__ is not None:
            args.__mbuddy__.free(ht.__slots__)
        args.__mbuddy__.free(t.CVoid(t.CUInt64T(ht), t.CPtr))
    args.__mbuddy__.free(t.CVoid(t.CUInt64T(args), t.CPtr))
