"""
内存操作模块，提供完全基于内存的Excel处理能力。
"""

import io
from pathlib import Path
from typing import Any, BinaryIO, Iterator

from .base import BaseMemoryWorkbook, WorkbookOperationsMixin
from .core.workbook import Workbook
from .core.worksheet import Worksheet
from .templates.template_workbook import TemplateWorkbook
from .utils.memory_manager import MemoryManager, MemoryStrategy


class InMemoryWorkbook(BaseMemoryWorkbook, WorkbookOperationsMixin):
    """完全基于内存的Excel工作簿类。"""

    def __init__(self, memory_manager: MemoryManager | None = None):
        """
        初始化内存工作簿。

        Args:
            memory_manager: 内存管理器实例
        """
        # 调用基类初始化,使用MEMORY_ONLY策略
        super().__init__(
            memory_manager=memory_manager,
            default_strategy=MemoryStrategy.MEMORY_ONLY,
        )
        self._workbook = Workbook(memory_manager=self._memory_manager)
        # 确保至少有一个工作表
        if len(self._workbook.sheetnames) == 0:
            self._workbook.create_worksheet("Sheet1")
        # 注册到内存管理器
        self._memory_manager.register_object(self)

    @classmethod
    def from_bytes(
        cls, data: bytes, memory_manager: MemoryManager | None = None
    ) -> "InMemoryWorkbook":
        """
        从字节数据创建内存工作簿。

        Args:
            data: Excel文件字节数据
            memory_manager: 内存管理器实例

        Returns:
            InMemoryWorkbook实例
        """
        instance = cls(memory_manager)

        # 使用BytesIO在内存中处理
        instance._workbook = Workbook.from_bytes(data)

        return instance

    @classmethod
    def from_template(
        cls,
        template_data: str | bytes | BinaryIO,
        context: dict[str, Any],
        memory_manager: MemoryManager | None = None,
    ) -> "InMemoryWorkbook":
        """
        从模板创建内存工作簿。

        Args:
            template_data: 模板数据（文件路径、字节数据或文件对象）
            context: 模板上下文变量
            memory_manager: 内存管理器实例

        Returns:
            InMemoryWorkbook实例
        """
        instance = cls(memory_manager)

        # 处理不同类型的模板数据
        if isinstance(template_data, str):
            # 文件路径
            template_workbook = TemplateWorkbook.load_template(template_data)
        elif isinstance(template_data, bytes):
            # 字节数据
            buffer = io.BytesIO(template_data)
            template_workbook = TemplateWorkbook.load_template(buffer)  # type: ignore[arg-type]
        else:
            # 文件对象
            template_workbook = TemplateWorkbook.load_template(template_data)  # type: ignore[arg-type]

        # 渲染模板
        rendered_workbook = template_workbook.render(context)

        # 将渲染后的工作簿转换为内存工作簿
        instance._workbook = rendered_workbook

        return instance

    def create_worksheet(self, name: str) -> Worksheet:
        """
        创建工作表。

        Args:
            name: 工作表名称

        Returns:
            工作表实例
        """
        return self._workbook.create_worksheet(name)

    def get_worksheet(self, name: str) -> Worksheet | None:
        """
        获取工作表。

        Args:
            name: 工作表名称

        Returns:
            工作表实例或None
        """
        return self._workbook.get_sheet_by_name(name)

    def delete_worksheet(self, name: str) -> bool:
        """
        删除工作表。

        Args:
            name: 工作表名称

        Returns:
            是否删除成功
        """
        return self._workbook.delete_worksheet(name)

    # get_sheet_names() 方法现在由 WorkbookOperationsMixin 提供

    def to_bytes(self) -> bytes:
        """
        转换为字节数据。

        Returns:
            Excel文件字节数据
        """
        return self._workbook.save_to_bytes()

    def to_file(self, file_path: str | Path) -> None:
        """
        保存到文件。

        Args:
            file_path: 文件路径
        """
        self._workbook.save(file_path)

    def to_stream(self) -> io.BytesIO:
        """
        转换为内存流。

        Returns:
            BytesIO对象
        """
        buffer = io.BytesIO()
        buffer.write(self.to_bytes())
        buffer.seek(0)
        return buffer

    def clone(self) -> "InMemoryWorkbook":
        """
        克隆内存工作簿。

        Returns:
            新的InMemoryWorkbook实例
        """
        # 获取当前工作簿的字节数据
        data = self.to_bytes()

        # 从字节数据创建新的内存工作簿
        return self.from_bytes(data, self._memory_manager)

    def merge_with(
        self, other: "InMemoryWorkbook", merge_strategy: str = "append"
    ) -> "InMemoryWorkbook":
        """
        合并另一个内存工作簿。

        Args:
            other: 另一个InMemoryWorkbook实例
            merge_strategy: 合并策略 ('append', 'replace', 'skip')

        Returns:
            新的合并后的InMemoryWorkbook实例
        """
        result = self.clone()

        for sheet_name in other.get_sheet_names():
            other_sheet = other.get_worksheet(sheet_name)
            if other_sheet is None:
                continue

            used_range = other_sheet.get_used_range()
            values = other_sheet.get_range_values(used_range)

            if merge_strategy == "append":
                # 追加工作表
                if sheet_name in result.get_sheet_names():
                    # 如果已存在，创建新名称
                    base_name = sheet_name
                    counter = 1
                    while f"{base_name}_{counter}" in result.get_sheet_names():
                        counter += 1
                    sheet_name = f"{base_name}_{counter}"

                new_sheet = result.create_worksheet(sheet_name)
                if values:
                    new_sheet.set_range_values("A1", values)

            elif merge_strategy == "replace":
                # 替换工作表
                if sheet_name in result.get_sheet_names():
                    result.delete_worksheet(sheet_name)
                target_sheet = result.create_worksheet(sheet_name)
                if values:
                    target_sheet.set_range_values("A1", values)

            elif merge_strategy == "skip":
                # 跳过已存在的工作表
                if sheet_name not in result.get_sheet_names():
                    target_sheet = result.create_worksheet(sheet_name)
                    if values:
                        target_sheet.set_range_values("A1", values)

        return result

    def get_memory_usage(self) -> dict[str, Any]:
        """
        获取内存使用情况。

        Returns:
            内存使用统计信息

        Note:
            扩展基类的get_memory_usage,添加工作簿特定的统计信息
        """
        # 调用基类方法获取基础统计
        stats = super().get_memory_usage()
        # 添加工作簿特定统计
        stats.update(
            {
                "total_size": len(self.to_bytes()),
                "sheet_count": len(self.get_sheet_names()),
            }
        )
        return stats

    # optimize_memory() 方法现在由 BaseMemoryWorkbook 提供

    # __enter__ 和 __exit__ 方法现在由 BaseMemoryWorkbook 提供

    def cleanup(self) -> None:
        """
        清理资源。

        Note:
            实现BaseMemoryWorkbook的抽象方法
        """
        # MemoryManager会自动处理弱引用清理,这里不需要额外操作
        pass


class InMemoryTemplateWorkbook(InMemoryWorkbook):
    """基于模板的内存工作簿类。"""

    def __init__(
        self,
        template_path: str | Path,
        memory_manager: MemoryManager | None = None,
    ):
        """
        初始化模板内存工作簿。

        Args:
            template_path: 模板文件路径
            memory_manager: 内存管理器实例
        """
        super().__init__(memory_manager)
        self._template_path = template_path
        self._template_workbook = TemplateWorkbook.load_template(template_path)

    def render(self, context: dict[str, Any]) -> "InMemoryTemplateWorkbook":
        """
        渲染模板。

        Args:
            context: 模板上下文变量

        Returns:
            渲染后的工作簿实例
        """
        rendered_workbook = self._template_workbook.render(context)
        self._workbook = rendered_workbook
        return self

    def render_worksheet(self, sheet_name: str, context: dict[str, Any]) -> Worksheet:
        """
        渲染特定工作表。

        Args:
            sheet_name: 工作表名称
            context: 模板上下文变量

        Returns:
            渲染后的工作表实例
        """
        return self._template_workbook.render_worksheet(sheet_name, context)

    def get_template_variables(self) -> set[str]:
        """
        获取模板变量。

        Returns:
            模板变量集合
        """
        return self._template_workbook.get_template_variables()


class MemoryOptimizedWorkbook(BaseMemoryWorkbook):
    """内存优化的工作簿类，适用于大数据量处理。"""

    def __init__(self, memory_manager: MemoryManager | None = None):
        """
        初始化内存优化工作簿。

        Args:
            memory_manager: 内存管理器实例
        """
        # 调用基类初始化,使用STREAMING策略
        super().__init__(
            memory_manager=memory_manager,
            default_strategy=MemoryStrategy.STREAMING,
        )
        self._workbook = Workbook(
            mode="write_only",
            memory_manager=self._memory_manager,
        )

    def create_sheet_with_streaming(
        self,
        sheet_name: str,
        data_iterator: Iterator[Any],
        *,
        headers: list[str] | None = None,
        batch_size: int = 1000,
    ) -> None:
        """
        使用流式处理创建工作表。

        Args:
            sheet_name: 工作表名称
            data_iterator: 数据迭代器
        """
        sheet = self._workbook.create_worksheet(sheet_name)

        if headers:
            sheet.append_row(headers)

        rows_since_flush = 0
        for normalized_row in self._iter_rows(data_iterator):
            if isinstance(normalized_row, dict):
                if headers:
                    sheet.append_row([normalized_row.get(h) for h in headers])
                else:
                    sheet.append_row(list(normalized_row.values()))
            elif isinstance(normalized_row, (list, tuple)):
                sheet.append_row(normalized_row)
            else:
                sheet.append_row([normalized_row])

            rows_since_flush += 1
            if rows_since_flush >= batch_size:
                self._memory_manager.optimize_memory()
                rows_since_flush = 0

        if rows_since_flush:
            self._memory_manager.optimize_memory()

    def to_bytes(self) -> bytes:
        """
        转换为字节数据。

        Returns:
            Excel文件字节数据
        """
        return self._workbook.save_to_memory()

    # __enter__ 和 __exit__ 方法现在由 BaseMemoryWorkbook 提供

    def cleanup(self) -> None:
        """
        清理临时文件和资源。

        Note:
            实现BaseMemoryWorkbook的抽象方法
        """
        self._workbook.close()

    def get_workbook(self) -> Workbook:
        """Expose underlying workbook for additional operations."""
        return self._workbook

    @staticmethod
    def _iter_rows(data_iterator: Iterator[Any]) -> Iterator[Any]:
        """Normalize incoming iterator into single rows."""
        for chunk in data_iterator:
            if chunk is None:
                continue
            if isinstance(chunk, (list, tuple)):
                if not chunk:
                    continue
                first = chunk[0]
                if isinstance(first, (list, tuple, dict)):
                    for row in chunk:
                        yield row
                else:
                    yield list(chunk)
            else:
                yield chunk


# 便捷函数
def create_in_memory_workbook() -> InMemoryWorkbook:
    """
    创建内存工作簿。

    Returns:
        InMemoryWorkbook实例
    """
    return InMemoryWorkbook()


def load_in_memory_workbook(data: bytes) -> InMemoryWorkbook:
    """
    从字节数据加载内存工作簿。

    Args:
        data: Excel文件字节数据

    Returns:
        InMemoryWorkbook实例
    """
    return InMemoryWorkbook.from_bytes(data)


def create_template_in_memory_workbook(
    template_data: str | bytes | BinaryIO, context: dict[str, Any]
) -> InMemoryWorkbook:
    """
    从模板创建内存工作簿。

    Args:
        template_data: 模板数据
        context: 模板上下文变量

    Returns:
        InMemoryWorkbook实例
    """
    return InMemoryWorkbook.from_template(template_data, context)
