from playwright.sync_api import Page, Locator
from typing import Dict, List, Optional


class TableHelper:
    """表格辅助工具类"""

    def __init__(self, page: Page):
        """初始化表格辅助工具

        Args:
            page: Playwright页面对象
        """
        self.page = page

    def get_row_count(self, table_selector: str, timeout: int = 10000) -> int:
        """获取表格行数

        Args:
            table_selector: 表格选择器
            timeout: 元素等待超时时间（毫秒）

        Returns:
            表格行数
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        rows = self.page.locator(f"{table_selector} tbody tr")
        return rows.count()

    def get_column_count(self, table_selector: str, timeout: int = 10000) -> int:
        """获取表格列数

        Args:
            table_selector: 表格选择器
            timeout: 元素等待超时时间（毫秒）

        Returns:
            表格列数
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        header_row = self.page.locator(f"{table_selector} thead tr")
        if header_row.count() > 0:
            cells = header_row.locator("th, td")
            return cells.count()

        first_row = self.page.locator(f"{table_selector} tbody tr:first-child")
        cells = first_row.locator("td")
        return cells.count()

    def get_cell_text(
        self, table_selector: str, row_index: int, col_index: int, timeout: int = 10000
    ) -> str:
        """获取单元格文本

        Args:
            table_selector: 表格选择器
            row_index: 行索引（从0开始）
            col_index: 列索引（从0开始）
            timeout: 元素等待超时时间（毫秒）

        Returns:
            单元格文本
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        cell = self.page.locator(
            f"{table_selector} tbody tr:nth-child({row_index + 1}) td:nth-child({col_index + 1})"
        )
        return cell.text_content()

    def get_row_data(self, table_selector: str, row_index: int, timeout: int = 10000) -> List[str]:
        """获取整行数据

        Args:
            table_selector: 表格选择器
            row_index: 行索引（从0开始）
            timeout: 元素等待超时时间（毫秒）

        Returns:
            行数据列表
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        row = self.page.locator(f"{table_selector} tbody tr:nth-child({row_index + 1})")
        cells = row.locator("td")

        row_data = []
        for i in range(cells.count()):
            row_data.append(cells.nth(i).text_content())

        return row_data

    def get_all_rows(self, table_selector: str, timeout: int = 10000) -> List[List[str]]:
        """获取所有行数据

        Args:
            table_selector: 表格选择器
            timeout: 元素等待超时时间（毫秒）

        Returns:
            所有行数据列表
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        rows = self.page.locator(f"{table_selector} tbody tr")

        all_rows = []
        for i in range(rows.count()):
            row_data = self.get_row_data(table_selector, i, timeout)
            all_rows.append(row_data)

        return all_rows

    def get_column_data(
        self, table_selector: str, col_index: int, timeout: int = 10000
    ) -> List[str]:
        """获取整列数据

        Args:
            table_selector: 表格选择器
            col_index: 列索引（从0开始）
            timeout: 元素等待超时时间（毫秒）

        Returns:
            列数据列表
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        rows = self.page.locator(f"{table_selector} tbody tr")

        column_data = []
        for i in range(rows.count()):
            cell_text = self.get_cell_text(table_selector, i, col_index, timeout)
            column_data.append(cell_text)

        return column_data

    def get_headers(self, table_selector: str, timeout: int = 10000) -> List[str]:
        """获取表头

        Args:
            table_selector: 表格选择器
            timeout: 元素等待超时时间（毫秒）

        Returns:
            表头列表
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        header_row = self.page.locator(f"{table_selector} thead tr")

        if header_row.count() > 0:
            headers = header_row.locator("th")
            header_list = []
            for i in range(headers.count()):
                header_list.append(headers.nth(i).text_content())
            return header_list

        return []

    def click_row(self, table_selector: str, row_index: int, timeout: int = 10000) -> None:
        """点击表格行

        Args:
            table_selector: 表格选择器
            row_index: 行索引（从0开始）
            timeout: 元素等待超时时间（毫秒）
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        row = self.page.locator(f"{table_selector} tbody tr:nth-child({row_index + 1})")
        row.click()

    def click_cell(
        self, table_selector: str, row_index: int, col_index: int, timeout: int = 10000
    ) -> None:
        """点击单元格

        Args:
            table_selector: 表格选择器
            row_index: 行索引（从0开始）
            col_index: 列索引（从0开始）
            timeout: 元素等待超时时间（毫秒）
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        cell = self.page.locator(
            f"{table_selector} tbody tr:nth-child({row_index + 1}) td:nth-child({col_index + 1})"
        )
        cell.click()

    def find_row_by_cell_text(
        self, table_selector: str, search_text: str, col_index: int = 0, timeout: int = 10000
    ) -> Optional[int]:
        """通过单元格文本查找行

        Args:
            table_selector: 表格选择器
            search_text: 要搜索的文本
            col_index: 搜索的列索引
            timeout: 元素等待超时时间（毫秒）

        Returns:
            找到的行索引，未找到则返回None
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        rows = self.page.locator(f"{table_selector} tbody tr")

        for i in range(rows.count()):
            cell_text = self.get_cell_text(table_selector, i, col_index, timeout)
            if cell_text and search_text in cell_text:
                return i

        return None

    def find_rows_by_cell_text(
        self, table_selector: str, search_text: str, col_index: int = 0, timeout: int = 10000
    ) -> List[int]:
        """通过单元格文本查找所有匹配行

        Args:
            table_selector: 表格选择器
            search_text: 要搜索的文本
            col_index: 搜索的列索引
            timeout: 元素等待超时时间（毫秒）

        Returns:
            找到的行索引列表
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        rows = self.page.locator(f"{table_selector} tbody tr")

        matching_rows = []
        for i in range(rows.count()):
            cell_text = self.get_cell_text(table_selector, i, col_index, timeout)
            if cell_text and search_text in cell_text:
                matching_rows.append(i)

        return matching_rows

    def sort_table_by_column(
        self, table_selector: str, col_index: int, ascending: bool = True, timeout: int = 10000
    ) -> None:
        """按列排序表格

        Args:
            table_selector: 表格选择器
            col_index: 列索引（从0开始）
            ascending: 是否升序排序
            timeout: 元素等待超时时间（毫秒）
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        header = self.page.locator(f"{table_selector} thead tr th:nth-child({col_index + 1})")
        header.click()

        if not ascending:
            header.click()

    def filter_table(self, table_selector: str, filter_text: str, timeout: int = 10000) -> None:
        """过滤表格

        Args:
            table_selector: 表格选择器
            filter_text: 过滤文本
            timeout: 元素等待超时时间（毫秒）
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")

        filter_input = self.page.locator(f"{table_selector} ~ .filter-input, .search-input")
        if filter_input.count() > 0:
            filter_input.fill(filter_text)
            self.page.keyboard.press("Enter")

    def select_row_checkbox(
        self, table_selector: str, row_index: int, timeout: int = 10000
    ) -> None:
        """选择行的复选框

        Args:
            table_selector: 表格选择器
            row_index: 行索引（从0开始）
            timeout: 元素等待超时时间（毫秒）
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        row = self.page.locator(f"{table_selector} tbody tr:nth-child({row_index + 1})")
        checkbox = row.locator("input[type='checkbox']")

        if checkbox.count() > 0:
            checkbox.check(force=True)

    def select_all_rows(self, table_selector: str, timeout: int = 10000) -> None:
        """选择所有行

        Args:
            table_selector: 表格选择器
            timeout: 元素等待超时时间（毫秒）
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")

        select_all_checkbox = self.page.locator(
            f"{table_selector} thead input[type='checkbox'], {table_selector} ~ .select-all-checkbox"
        )

        if select_all_checkbox.count() > 0:
            select_all_checkbox.check(force=True)

    def deselect_all_rows(self, table_selector: str, timeout: int = 10000) -> None:
        """取消选择所有行

        Args:
            table_selector: 表格选择器
            timeout: 元素等待超时时间（毫秒）
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")

        select_all_checkbox = self.page.locator(
            f"{table_selector} thead input[type='checkbox'], {table_selector} ~ .select-all-checkbox"
        )

        if select_all_checkbox.count() > 0:
            select_all_checkbox.uncheck(force=True)

    def get_row_by_id(
        self, table_selector: str, row_id: str, timeout: int = 10000
    ) -> Optional[Locator]:
        """通过ID获取行

        Args:
            table_selector: 表格选择器
            row_id: 行ID
            timeout: 元素等待超时时间（毫秒）

        Returns:
            行元素，未找到则返回None
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        row = self.page.locator(f"{table_selector} tbody tr[data-id='{row_id}']")

        if row.count() > 0:
            return row

        return None

    def click_row_action_button(
        self, table_selector: str, row_index: int, action: str = "edit", timeout: int = 10000
    ) -> None:
        """点击行操作按钮

        Args:
            table_selector: 表格选择器
            row_index: 行索引（从0开始）
            action: 操作类型（edit/delete/view等）
            timeout: 元素等待超时时间（毫秒）
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        row = self.page.locator(f"{table_selector} tbody tr:nth-child({row_index + 1})")

        action_button = row.locator(f"button.{action}, .{action}-button")
        if action_button.count() > 0:
            action_button.click()
        else:
            actions_cell = row.locator("td:last-child")
            actions_cell.click()

    def wait_for_table_load(
        self, table_selector: str, timeout: int = 10000
    ) -> None:
        """等待表格加载完成

        Args:
            table_selector: 表格选择器
            timeout: 等待超时时间（毫秒）
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        self.page.wait_for_timeout(500)

    def is_table_empty(self, table_selector: str, timeout: int = 10000) -> bool:
        """检查表格是否为空

        Args:
            table_selector: 表格选择器
            timeout: 元素等待超时时间（毫秒）

        Returns:
            表格是否为空
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")

        empty_message = self.page.locator(
            f"{table_selector} ~ .no-data, {table_selector} ~ .empty-message"
        )

        if empty_message.count() > 0 and empty_message.is_visible():
            return True

        return self.get_row_count(table_selector, timeout) == 0

    def get_table_data_as_dict(
        self, table_selector: str, timeout: int = 10000
    ) -> List[Dict[str, str]]:
        """获取表格数据为字典列表

        Args:
            table_selector: 表格选择器
            timeout: 元素等待超时时间（毫秒）

        Returns:
            表格数据字典列表，键为列名，值为单元格内容
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")

        headers = self.get_headers(table_selector, timeout)
        all_rows = self.get_all_rows(table_selector, timeout)

        table_data = []
        for row_data in all_rows:
            row_dict = {}
            for i, header in enumerate(headers):
                if i < len(row_data):
                    row_dict[header] = row_data[i]
            table_data.append(row_dict)

        return table_data

    def scroll_to_row(self, table_selector: str, row_index: int, timeout: int = 10000) -> None:
        """滚动到指定行

        Args:
            table_selector: 表格选择器
            row_index: 行索引（从0开始）
            timeout: 元素等待超时时间（毫秒）
        """
        self.page.wait_for_selector(table_selector, timeout=timeout, state="visible")
        row = self.page.locator(f"{table_selector} tbody tr:nth-child({row_index + 1})")
        row.scroll_into_view_if_needed()