import typing
import contextlib
import importlib.metadata
import io
import re
import sys
import warnings
from typing import Any, Callable, Literal, Optional, Tuple


# 辅助函数 get_dependency_info 保持不变
def get_dependency_info(obj: Any) -> Tuple[Optional[str], Optional[str]]:
    """
    尝试查找一个对象的类型所依赖的包名和版本号。
    """
    try:
        module_name = type(obj).__module__
        # 跳过内建类型和主执行脚本中定义的类型
        if module_name == 'builtins' or module_name == '__main__':
            return (None, None)

        package_name = module_name.split('.')[0]
        version = importlib.metadata.version(package_name)
        return (package_name, version)
    except importlib.metadata.PackageNotFoundError:
        return (None, None)
    except Exception:
        return (None, None)


def type_summary(obj: Any, inspect_mode: bool = False) -> str:
    """返回对象的类型摘要

    此函数接受任何类型的对象，并返回一个字符串，描述该对象的类型和一些基本信息。

    Examples
    --------
    ```py
    >>> type_summary("hello")
    'hello'
    >>> type_summary(42)
    '42'
    >>> type_summary(3.14159)
    '3.142'
    >>> type_summary([1, 2, 3])
    'int[3]'
    >>> type_summary({'a': 1, 'b': 2})
    "dict[str: int]"
    >>> import numpy as np
    >>> type_summary(np.array([1, 2, 3]))
    'ndarray@int64(3,)'
    ```
    """
    summary = ''
    # 1. 生成基础的类型摘要
    if isinstance(obj, str):
        summary = f'str[{len(obj)}]'
    elif isinstance(obj, (int, bool)):
        summary = str(obj)
    elif isinstance(obj, float):
        summary = f'{obj:.3f}'
    elif isinstance(obj, (list, tuple)):
        if not obj:
            summary = f'{type(obj).__name__}[]'
        else:
            item_type_name = type(obj[0]).__name__
            summary = f'{type(obj).__name__}[{len(obj)}] of {item_type_name}'
    elif isinstance(obj, dict):
        if not obj:
            summary = f'{type(obj).__name__}[]'
        else:
            key_type_name = type(next(iter(obj.keys()))).__name__
            value_type_name = type(next(iter(obj.values()))).__name__
            summary = f'{type(obj).__name__}[{key_type_name}: {value_type_name}]'
    elif hasattr(obj, 'shape') and hasattr(obj, 'dtype'):
        summary = f'{type(obj).__name__}@{obj.dtype}{obj.shape}'
    else:
        summary = type(obj).__name__

    # 2. 如果开启了检查模式，则尝试附加依赖信息
    if inspect_mode:
        package, version = get_dependency_info(obj)
        if package and version:
            # 只有在找到包信息时才附加
            return f'{summary} (from {package}=={version})'

    # 默认或未找到依赖信息时，返回基础摘要
    return summary


COLOR_MAP = {
    'red': '\033[91m',
    'green': '\033[92m',
    'yellow': '\033[93m',
    'blue': '\033[94m',
    'magenta': '\033[95m',
    'cyan': '\033[96m',
    'white': '\033[97m',
    'gray': '\033[90m',
    'reset': '\033[0m',
}


COLOR_TYPE = Literal['red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white', 'gray']


def colored(text: str, color: COLOR_TYPE):
    """为文本添加颜色

    只有在 TTY 环境下才会添加颜色代码。

    Parameters
    ----------
    - `text` : `str`
        - 原始文本
    - `color` : `COLOR_TYPE`
        - 颜色类型
        - `"red"` | `"green"` | `"yellow"` | `"blue"` | `"magenta"` | `"cyan"` | `"white"` | `"gray"`
    """
    if sys.stdout.isatty():
        return f"{COLOR_MAP[color]}{text}{COLOR_MAP['reset']}"
    return text


def print_colored(color: COLOR_TYPE, *args, **kwargs):
    """打印带颜色的文本

    Parameters
    ----------
    - `color` : `COLOR_TYPE`
        - 颜色类型
        - `"red"` | `"green"` | `"yellow"` | `"blue"` | `"magenta"` | `"cyan"` | `"white"` | `"gray"`
    """

    total_text = ' '.join([str(arg) for arg in args])
    print(colored(total_text, color), **kwargs)


def print_debug(*args, **kwargs):
    """打印调试信息（白色）"""
    print_colored('gray', *args, **kwargs)


def print_info(*args, **kwargs):
    """打印信息（青色）"""
    print_colored('cyan', *args, **kwargs)


def print_warning(*args, **kwargs):
    """打印警告信息（黄色）"""
    print_colored('yellow', *args, **kwargs)


print_warn = print_warning  # alias


def print_table(
    data: typing.Sequence[typing.Sequence[Any]]
    | typing.Sequence[typing.Mapping[str, Any]],
    header: Optional[typing.Sequence[str]] = None,
    fn_fmt: Callable[..., str] | None = None,
    color: Literal['red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white']
    | None = None,
    index: bool = False,
):
    """打印表格

    支持两种数据格式：
    - `list[list[Any]]`: 传统的二维列表格式
    - `list[dict[str, Any]]`: 字典列表格式，会自动提取 header 和转换数据

    Parameters
    ----------
    - `data` : `list[list[Any]] | list[dict[str, Any]]`
        - 表格数据，支持二维列表或字典列表
    - `header` : `list[str]`, optional
        - 表头，如果数据是字典列表且未提供 header，会自动从第一个字典的 keys 提取
    - `fn_fmt` : `(item) -> str | (key, item) -> str`, optional
        - 格式化函数，用于自定义每个单元格的显示格式。
        - 支持 `fn(value)` 和 `fn(key, value)` 两种形式。
    - `color` : `COLOR_TYPE`, optional
        - 表格颜色
    - `index` : `bool`, optional
        - 是否显示行索引, by default False
    """
    import inspect

    from prettytable import PrettyTable

    # from zyplib.utils.ensure import ensure_scalar
    from zyplib.utils.kv import flatten_dict

    # 检查 fn_fmt 的参数数量
    num_fn_fmt_args = 0
    if fn_fmt:
        try:
            num_fn_fmt_args = len(inspect.signature(fn_fmt).parameters)
        except (ValueError, TypeError):
            num_fn_fmt_args = 1  # 默认值为 1

    # 检测数据格式并处理
    if data and isinstance(data[0], dict):
        first_row = flatten_dict(typing.cast(typing.Mapping[str, Any], data[0]))
        original_keys = list(first_row.keys())
        if header is not None:
            import warnings

            warnings.warn(
                '当数据为字典列表格式时，会自动从字典键提取表头，用户指定的 header 将被忽略',
                UserWarning,
                stacklevel=2,
            )
        header = original_keys
        converted_data = []
        for row_dict in data:
            row_dict = flatten_dict(typing.cast(typing.Mapping[str, Any], row_dict))
            row = [row_dict.get(key, None) for key in original_keys]
            converted_data.append(row)
        data = converted_data

    table = PrettyTable()
    if header:
        if index:
            table.field_names = ['Index'] + typing.cast(list[str], header)
        else:
            table.field_names = typing.cast(typing.Any, header)

    for i, raw_row in enumerate(data):
        row = []
        for j, item in enumerate(raw_row):
            if fn_fmt:
                if num_fn_fmt_args == 2:
                    key = header[j] if header else str(j)
                    val = fn_fmt(key, item)
                else:
                    val = fn_fmt(item)
            else:
                val = str(item)
            row.append(val)
        if index:
            table.add_row([i] + row)
        else:
            table.add_row(row)

    if not color:
        print(table)
    else:
        from zyplib.utils.print import print_colored

        print_colored(color, table.get_string())


def warning(
    message: str,
    category: type[Warning] | None = None,
    stacklevel: int = 2,
    source: Any | None = None,
):
    """打印警告信息（黄色）等价于调用 `warnings.warn` 但是增加颜色"""
    warnings.warn(colored(message, 'yellow'), category, stacklevel, source)


def print_error(*args, **kwargs):
    """打印错误信息（红色）"""
    print_colored('red', *args, **kwargs)


class PrintCapture:
    """捕获 print 输出并同时保存到文件的上下文管理器

    在保持终端正常显示的同时，将所有 print 输出保存到指定文件中。
    写入文件时会自动清理 ANSI 颜色代码。

    Examples
    --------
    ```python
    # 基本用法
    with PrintCapture("log.txt"):
        print("这会同时显示在终端和保存到文件")
        print_info("带颜色的信息")

    # 追加模式
    with PrintCapture("log.txt", mode='a'):
        print("追加到现有文件")
    """

    def __init__(
        self,
        filename: str,
        mode: Literal['w', 'a'] = 'w',
        encoding: str = 'utf-8',
        flush_interval: int = 10,
    ):
        """初始化 PrintCapture

        Parameters
        ----------
        - `filename` : `str`
            - 输出文件名
        - `mode` : `Literal['w', 'a']`, optional
            - 文件写入模式，'w' 为覆盖写入，'a' 为追加写入, by default 'w'
        - `encoding` : `str`, optional
            - 文件编码格式, by default 'utf-8'
        - `flush_interval` : `int`, optional
            - 刷新间隔，每写入多少次后才刷新到磁盘，1表示立即刷新, by default 1
        - `auto_detect_ansi` : `bool`, optional
            - 是否自动检测ANSI代码，只有检测到时才进行正则替换, by default True
        """
        self.filename = filename
        self.mode = mode
        self.encoding = encoding
        self.flush_interval = flush_interval

        self.file = None
        self.original_stdout = None
        self.original_stderr = None

        # 用于清理 ANSI 颜色代码的正则表达式
        self.ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')

    def __enter__(self):
        """进入上下文，开始捕获输出"""
        # 打开输出文件
        self.file = open(self.filename, self.mode, encoding=self.encoding)

        # 保存原始的 stdout 和 stderr
        self.original_stdout = sys.stdout
        self.original_stderr = sys.stderr

        # 创建自定义的输出流
        sys.stdout = self._create_tee_stream(self.original_stdout)
        sys.stderr = self._create_tee_stream(self.original_stderr)

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """退出上下文，恢复原始输出流"""
        # 最终刷新：确保所有缓冲的内容都写入磁盘
        if self.file:
            self.file.flush()

        # 恢复原始输出流
        sys.stdout = self.original_stdout
        sys.stderr = self.original_stderr

        # 关闭文件
        if self.file:
            self.file.close()

    def _create_tee_stream(self, original_stream):
        """创建一个同时输出到终端和文件的流对象"""

        class TeeStream:
            def __init__(tee_self, original, file_handle, ansi_cleaner, flush_interval):
                tee_self.original = original
                tee_self.file = file_handle
                tee_self.ansi_cleaner = ansi_cleaner
                tee_self.flush_interval = flush_interval
                tee_self.write_count = 0  # 写入计数器

            def write(tee_self, text):
                # 写入原始流（终端）- 保持颜色
                tee_self.original.write(text)

                # 写入文件 - 智能处理颜色代码
                if tee_self.file:
                    if '\x1b' in text:
                        clean_text = tee_self.ansi_cleaner.sub('', text)
                    else:
                        clean_text = text

                    tee_self.file.write(clean_text)

                    # 性能优化：控制刷新频率
                    tee_self.write_count += 1
                    if tee_self.write_count >= tee_self.flush_interval:
                        tee_self.file.flush()
                        tee_self.write_count = 0

                return len(text)

            def flush(tee_self):
                tee_self.original.flush()
                if tee_self.file:
                    tee_self.file.flush()
                    tee_self.write_count = 0  # 重置计数器

            def __getattr__(tee_self, name):
                # 将其他属性委托给原始流
                return getattr(tee_self.original, name)

        return TeeStream(
            original_stream, self.file, self.ansi_escape, self.flush_interval
        )


@contextlib.contextmanager
def capture_print(
    filename: str,
    mode: Literal['w', 'a'] = 'w',
    encoding: str = 'utf-8',
    flush_interval: int = 10,
):
    """捕获 print 输出并同时保存到文件

    这是 PrintCapture 的便捷函数形式。在保持终端正常显示的同时，
    将所有 print 输出保存到指定文件中。写入文件时会自动清理 ANSI 颜色代码。

    Parameters
    ----------
    - `filename` : `str`
        - 输出文件名
    - `mode` : `Literal['w', 'a']`, optional
        - 文件写入模式，'w' 为覆盖写入，'a' 为追加写入, by default 'w'
    - `encoding` : `str`, optional
        - 文件编码格式, by default 'utf-8'
    - `flush_interval` : `int`, optional
        - 刷新间隔，每写入多少次后才刷新到磁盘，1表示立即刷新, by default 1
    - `auto_detect_ansi` : `bool`, optional
        - 是否自动检测ANSI代码，只有检测到时才进行正则替换, by default True

    Examples
    --------
    ```python
    # 基本用法
    with capture_print("training_log.txt"):
        print("训练开始...")
        print_info("这是一条带颜色的信息")
        print_warning("这是一条警告")

    # 追加模式
    with capture_print("training_log.txt", mode='a'):
        print("继续训练...")

    # 高性能模式：减少刷新频率，提高性能
    with capture_print("training_session.log", flush_interval=50):
        for epoch in range(1000):
            print(f"Epoch {epoch+1}/1000")
            print_info(f"Loss: {0.5 * (1-epoch/1000):.4f}")
    ```
    """
    with PrintCapture(
        filename, mode=mode, encoding=encoding, flush_interval=flush_interval
    ) as capture:
        yield capture
