import t, c
from stdint import *
import string
import memhub
import stdio
import hashtable


# ============================================================
# argparse.py — 命令行参数解析库（Python argparse 风格）
#
# 用法:
#   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)
# ============================================================

# 类型常量
STRING: t.CDefine = 0
INT: t.CDefine = 1
BOOL: t.CDefine = 2

# 动作常量
STORE: t.CDefine = 0
STORE_TRUE: t.CDefine = 1
STORE_FALSE: t.CDefine = 2
COUNT: t.CDefine = 3

# 限制
MAX_ARGS: t.CDefine = 32

# 模块级 mbuddy 指针（由用户在 main 中赋值）
_mbuddy: memhub.MemBuddy | t.CPtr


# ============================================================
# Argument — 参数定义结构体
# ============================================================

class Argument:
    name: str               # "--input" 或 "output"（位置参数）
    short_name: str         # "-i" 或 None
    help_text: str          # 帮助文本
    arg_type: INT           # STRING, INT, BOOL
    action: INT             # STORE, STORE_TRUE, STORE_FALSE, COUNT
    default_int: INT        # INT 类型默认值
    default_str: str        # STRING 类型默认值
    required: bool          # 是否必须
    is_positional: bool     # 是否位置参数
    dest: str               # 查找键名（"--input" → "input"）


# ============================================================
# ParsedArgs — 解析结果
# ============================================================

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:
        """获取字符串值。键不存在返回 None。"""
        if self._ht.__contains__(name) == 0: return None
        return self._ht[name]

    def get_int(self, name: str) -> INT:
        """获取整数值。键不存在返回 0。"""
        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) -> INT:
        """获取布尔值（返回 1 或 0）。键不存在返回 0。"""
        if self._ht.__contains__(name) == 0: return 0
        ptr: t.CPtr = self._ht[name]
        return 1 if c.Deref(t.CInt(t.CUInt64T(ptr), t.CPtr)) != 0 else 0

    def has(self, name: str) -> bool:
        """检查键是否存在。"""
        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 — 命令行参数解析器
# ============================================================

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:
        mb: memhub.MemBuddy | t.CPtr = pool
        if mb is None: mb = _mbuddy
        return mb.alloc(40)

    def __init__(self, prog: str = "", description: str = "",
                 pool: memhub.MemBuddy | t.CPtr = None):
        mb: memhub.MemBuddy | t.CPtr = pool
        if mb is None: mb = _mbuddy
        self.__mbuddy__ = mb
        self._prog = prog
        self._description = description
        self._arg_count = 0
        self._args = mb.alloc(MAX_ARGS * Argument.__sizeof__())

    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):
        """注册一个命令行参数。

        Args:
            name: 参数名（"--input" 或 "output"）
            short: 短选项（"-i"），可选
            arg_type: 类型（STRING/INT/BOOL）
            default: INT 默认值
            default_str: STRING 默认值
            required: 是否必须
            action: 动作（STORE/STORE_TRUE/STORE_FALSE/COUNT）
            help: 帮助文本
        """
        if self._arg_count >= MAX_ARGS: return
        idx: INT = self._arg_count
        arg: Argument | t.CPtr = t.CPtr(t.CUInt64T(self._args) + idx * Argument.__sizeof__())
        arg.name = 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
        # 判断是否位置参数
        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
        self._arg_count = idx + 1

    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:
        """精确匹配长选项名"""
        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:
        """匹配 --name=value 中的 name 前缀"""
        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)
            if v != 0:
                args._set_bool(arg.dest, True)
            else:
                args._set_bool(arg.dest, False)

    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:
        """解析命令行参数。

        Args:
            argc: 参数个数
            argv: 参数数组（char**，argv[0] 为程序名）
        Returns:
            ParsedArgs 指针，失败返回 None
        """
        if self.__mbuddy__ is None: return None

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

        # 初始化默认值
        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

        # 解析 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
        name_len: INT = 0

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

            # "--" 结束符：后续全部为位置参数
            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

            # 长选项 --name 或 --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
                    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（值在下一个 argv 或标志）
                    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

            # 短选项 -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

            # 位置参数
            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

        # 检查必须参数
        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):
        """打印帮助信息"""
        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):
        """打印用法简述"""
        stdio.printf("usage: %s [options]\n", self._prog)


# ============================================================
# release — 释放 ParsedArgs 及其内部资源
# ============================================================

def release(args: ParsedArgs | t.CPtr):
    """释放 ParsedArgs 对象及其内部 HashTable。"""
    if args is None: return
    if args.__mbuddy__ is None: return
    if args._ht is not None:
        if args._ht.__slots__ is not None:
            args.__mbuddy__.free(args._ht.__slots__)
        args.__mbuddy__.free(t.CVoid(t.CUInt64T(args._ht), t.CPtr))
    args.__mbuddy__.free(t.CVoid(t.CUInt64T(args), t.CPtr))