#!/usr/bin/env python3
"""
飞书文档通用 I/O 模块
提供干净的读写 API，供各专项技能调用。

支持：
  - 飞书文档表格（docx/wiki 中的 table block）
  - 飞书多维表格（bitable）

用法（作为模块被其他技能 import）：
  from feishu_io import FeishuDocIO
  io = FeishuDocIO(app_id, app_secret)
  rows = io.read_table_column(url, "中文台词")
  io.write_table_column(url, "英文翻译", results)
"""
import requests
import json
import re
import time
import os
from typing import List, Dict, Optional, Tuple


class FeishuDocIO:
    """飞书文档通用读写能力"""

    def __init__(self, app_id: str = None, app_secret: str = None):
        self.app_id = app_id or os.getenv("FEISHU_APP_ID", "")
        self.app_secret = app_secret or os.getenv("FEISHU_APP_SECRET", "")
        self._token = None
        self._token_url = "https://open.feishu.cn/open-apis/auth/v3/tenant_access_token/internal"

    # ============================================================
    # 凭证管理
    # ============================================================
    def _get_token(self) -> str:
        """获取 tenant_access_token（自动缓存）"""
        if self._token:
            return self._token
        if not self.app_id or not self.app_secret:
            raise ValueError("缺少飞书应用凭证，请设置 FEISHU_APP_ID 和 FEISHU_APP_SECRET 环境变量")
        resp = requests.post(self._token_url,
                             json={"app_id": self.app_id, "app_secret": self.app_secret})
        data = resp.json()
        if "tenant_access_token" not in data:
            raise RuntimeError(f"获取飞书凭证失败: {data}")
        self._token = data["tenant_access_token"]
        return self._token

    def _headers(self) -> dict:
        return {"Authorization": f"Bearer {self._get_token()}", "Content-Type": "application/json"}

    # ============================================================
    # URL 解析
    # ============================================================
    @staticmethod
    def parse_url(url: str) -> Dict:
        """解析飞书链接，返回文档类型和关键 ID"""
        result = {"type": "unknown"}
        if "/wiki/" in url:
            result["type"] = "wiki"
            result["wiki_token"] = re.findall(r"/wiki/([a-zA-Z0-9]+)", url)[0]
        elif "/docx/" in url:
            result["type"] = "docx"
            result["doc_id"] = re.findall(r"/docx/([a-zA-Z0-9]+)", url)[0]
        elif "/bitable/" in url:
            result["type"] = "bitable"
            result["app_token"] = re.findall(r"/bitable/([a-zA-Z0-9]+)", url)[0]
            table_match = re.findall(r"table=([a-zA-Z0-9]+)", url)
            if table_match:
                result["table_id"] = table_match[0]
        return result

    def _resolve_wiki_to_doc_id(self, wiki_token: str) -> str:
        """将 wiki token 转换为实际 doc_id"""
        resp = requests.get(
            f"https://open.feishu.cn/open-apis/wiki/v2/spaces/get_node?token={wiki_token}",
            headers=self._headers())
        data = resp.json()
        if data.get("code") != 0:
            raise RuntimeError(f"获取wiki节点失败: {data}")
        return data["data"]["node"]["obj_token"]

    # ============================================================
    # 文档表格：读取
    # ============================================================
    def _find_tables(self, doc_id: str) -> List[Dict]:
        """获取文档中所有表格的基本信息"""
        resp = requests.get(
            f"https://open.feishu.cn/open-apis/docx/v1/documents/{doc_id}/blocks?page_size=500",
            headers=self._headers())
        data = resp.json()
        if data.get("code") != 0:
            raise RuntimeError(f"获取文档blocks失败: {data}")
        tables = []
        for b in data["data"]["items"]:
            if b.get("block_type") == 31:
                prop = b.get("table", {}).get("property", {})
                tables.append({
                    "block_id": b["block_id"],
                    "rows": prop.get("row_size", 0),
                    "cols": prop.get("column_size", 0)
                })
        return tables

    def _get_table_cells(self, doc_id: str, table_block_id: str) -> List[Dict]:
        """获取表格所有单元格 block"""
        resp = requests.get(
            f"https://open.feishu.cn/open-apis/docx/v1/documents/{doc_id}/blocks/{table_block_id}/children?page_size=500",
            headers=self._headers())
        data = resp.json()
        if data.get("code") != 0:
            raise RuntimeError(f"获取表格单元格失败: {data}")
        return data["data"]["items"]

    def _read_cell_text(self, doc_id: str, cell_block_id: str) -> str:
        """读取单个单元格的文本内容，自动处理换行和HTML标签"""
        resp = requests.get(
            f"https://open.feishu.cn/open-apis/docx/v1/documents/{doc_id}/blocks/{cell_block_id}/children?page_size=50",
            headers=self._headers())
        data = resp.json()
        if data.get("code") != 0:
            return ""
        paragraphs = []
        for child in data["data"].get("items", []):
            if child.get("block_type") == 2:
                # 同一text block内的text_run拼接成一段（它们是同一段落的不同片段）
                parts = []
                for el in child.get("text", {}).get("elements", []):
                    content = el.get("text_run", {}).get("content", "")
                    if content:
                        parts.append(content)
                line = "".join(parts).strip()
                if line:
                    paragraphs.append(line)
        result = "\n".join(paragraphs)
        # 清洗HTML换行标签 → 真实换行
        result = re.sub(r'<br\s*/?\s*>', '\n', result, flags=re.IGNORECASE)
        # 清理连续空行
        result = re.sub(r'\n{3,}', '\n\n', result)
        return result.strip()

    def read_table_column(self, url: str, column_name: str,
                          table_index: int = 0, start_row: int = 1,
                          end_row: int = None) -> List[Dict]:
        """
        从飞书文档表格中读取指定列的数据。

        Args:
            url: 飞书文档链接
            column_name: 目标列的表头名称
            table_index: 表格序号（文档中第几个表格，从0开始）
            start_row: 起始数据行（0为表头行，1为第一行数据）
            end_row: 结束行（None=读到最后）

        Returns:
            [{"row": 行号, "content": "单元格文本"}, ...]
        """
        url_info = self.parse_url(url)
        if url_info["type"] == "wiki":
            doc_id = self._resolve_wiki_to_doc_id(url_info["wiki_token"])
        elif url_info["type"] == "docx":
            doc_id = url_info["doc_id"]
        else:
            raise ValueError(f"read_table_column 仅支持文档表格(docx/wiki)，当前类型: {url_info['type']}")

        # 找到目标表格
        tables = self._find_tables(doc_id)
        if table_index >= len(tables):
            raise ValueError(f"文档中只有 {len(tables)} 个表格，无法访问第 {table_index} 个")
        table = tables[table_index]
        total_cols = table["cols"]

        # 获取所有单元格
        cells = self._get_table_cells(doc_id, table["block_id"])

        # 读取表头行，找到目标列的索引
        header_cells = cells[:total_cols]
        col_index = None
        for i, cell in enumerate(header_cells):
            header_text = self._read_cell_text(doc_id, cell["block_id"])
            if column_name in header_text:
                col_index = i
                break
        if col_index is None:
            available = [self._read_cell_text(doc_id, c["block_id"]) for c in header_cells]
            raise ValueError(f"未找到列 '{column_name}'，可用列: {available}")

        # 读取数据行
        total_rows = len(cells) // total_cols
        if end_row is None:
            end_row = total_rows - 1
        end_row = min(end_row, total_rows - 1)

        results = []
        for row in range(start_row, end_row + 1):
            cell_idx = row * total_cols + col_index
            if cell_idx < len(cells):
                content = self._read_cell_text(doc_id, cells[cell_idx]["block_id"])
                results.append({"row": row, "content": content})
                time.sleep(0.1)  # 避免API频率限制
        return results

    def read_table_columns(self, url: str, column_names: List[str],
                           table_index: int = 0, start_row: int = 1,
                           end_row: int = None) -> List[Dict]:
        """
        从飞书表格中读取多列数据。

        Returns:
            [{"row": 行号, "col1_name": "内容", "col2_name": "内容"}, ...]
        """
        url_info = self.parse_url(url)
        if url_info["type"] == "wiki":
            doc_id = self._resolve_wiki_to_doc_id(url_info["wiki_token"])
        elif url_info["type"] == "docx":
            doc_id = url_info["doc_id"]
        else:
            raise ValueError(f"仅支持文档表格(docx/wiki)")

        tables = self._find_tables(doc_id)
        table = tables[table_index]
        total_cols = table["cols"]
        cells = self._get_table_cells(doc_id, table["block_id"])

        # 解析表头
        header_cells = cells[:total_cols]
        col_map = {}  # column_name → col_index
        for i, cell in enumerate(header_cells):
            header_text = self._read_cell_text(doc_id, cell["block_id"])
            for name in column_names:
                if name in header_text:
                    col_map[name] = i

        missing = [n for n in column_names if n not in col_map]
        if missing:
            raise ValueError(f"未找到列: {missing}")

        total_rows = len(cells) // total_cols
        if end_row is None:
            end_row = total_rows - 1
        end_row = min(end_row, total_rows - 1)

        results = []
        for row in range(start_row, end_row + 1):
            row_data = {"row": row}
            for name, col_idx in col_map.items():
                cell_idx = row * total_cols + col_idx
                if cell_idx < len(cells):
                    row_data[name] = self._read_cell_text(doc_id, cells[cell_idx]["block_id"])
            results.append(row_data)
            time.sleep(0.1)
        return results

    # ============================================================
    # 文档表格：写入
    # ============================================================
    def _cell_has_content(self, doc_id: str, cell_id: str) -> Tuple[bool, List]:
        """检查单元格是否已有文本内容"""
        resp = requests.get(
            f"https://open.feishu.cn/open-apis/docx/v1/documents/{doc_id}/blocks/{cell_id}/children?page_size=50",
            headers=self._headers())
        data = resp.json()
        if data.get("code") != 0:
            return False, []
        children = data["data"].get("items", [])
        for child in children:
            if child.get("block_type") == 2:
                for el in child.get("text", {}).get("elements", []):
                    if el.get("text_run", {}).get("content", "").strip():
                        return True, children
        return False, children

    def _write_cell(self, doc_id: str, cell_id: str, text: str) -> bool:
        """写入文本到单个单元格"""
        # 清理现有子blocks
        _, existing = self._cell_has_content(doc_id, cell_id)
        for child in existing:
            requests.delete(
                f"https://open.feishu.cn/open-apis/docx/v1/documents/{doc_id}/blocks/{child['block_id']}",
                headers=self._headers())

        # 写入新内容
        text_elements = [{"text_run": {"content": line + "\n"}} for line in text.split("\n")]
        payload = {
            "children": [{"block_type": 2, "text": {"elements": text_elements, "style": {}}}],
            "index": 0
        }
        resp = requests.post(
            f"https://open.feishu.cn/open-apis/docx/v1/documents/{doc_id}/blocks/{cell_id}/children",
            headers=self._headers(), json=payload)
        return resp.json().get("code") == 0

    def write_table_column(self, url: str, column_name: str, data: List[Dict],
                           table_index: int = 0, skip_existing: bool = True,
                           batch_size: int = 10, batch_delay: float = 0.3) -> Dict:
        """
        向飞书文档表格指定列写入数据。

        Args:
            url: 飞书文档链接
            column_name: 目标列的表头名称
            data: [{"row": 行号, "content": "要写入的文本"}, ...]
            table_index: 表格序号
            skip_existing: 是否跳过已有内容的单元格
            batch_size: 每批处理行数
            batch_delay: 每次写入间隔（秒）

        Returns:
            {"success": 成功数, "skipped": 跳过数, "failed": 失败数}
        """
        url_info = self.parse_url(url)
        if url_info["type"] == "wiki":
            doc_id = self._resolve_wiki_to_doc_id(url_info["wiki_token"])
        elif url_info["type"] == "docx":
            doc_id = url_info["doc_id"]
        else:
            raise ValueError(f"write_table_column 仅支持文档表格(docx/wiki)")

        tables = self._find_tables(doc_id)
        table = tables[table_index]
        total_cols = table["cols"]
        cells = self._get_table_cells(doc_id, table["block_id"])

        # 找到目标列
        header_cells = cells[:total_cols]
        col_index = None
        for i, cell in enumerate(header_cells):
            header_text = self._read_cell_text(doc_id, cell["block_id"])
            if column_name in header_text:
                col_index = i
                break
        if col_index is None:
            raise ValueError(f"未找到列 '{column_name}'")

        stats = {"success": 0, "skipped": 0, "failed": 0}

        for item in data:
            row = item["row"]
            text = item["content"]
            cell_idx = row * total_cols + col_index
            if cell_idx >= len(cells):
                stats["failed"] += 1
                continue

            cell_id = cells[cell_idx]["block_id"]

            # 跳过已有内容
            if skip_existing:
                has_content, _ = self._cell_has_content(doc_id, cell_id)
                if has_content:
                    stats["skipped"] += 1
                    print(f"  行{row}: 跳过（已有内容）")
                    continue

            if self._write_cell(doc_id, cell_id, text):
                stats["success"] += 1
                print(f"  行{row}: ✅ 写入成功")
            else:
                stats["failed"] += 1
                print(f"  行{row}: ❌ 写入失败")

            time.sleep(batch_delay)

        return stats

    # ============================================================
    # 多维表格：读写
    # ============================================================
    def read_bitable_column(self, url: str, column_name: str) -> List[Dict]:
        """从多维表格读取指定列"""
        url_info = self.parse_url(url)
        if url_info["type"] != "bitable":
            raise ValueError("仅支持多维表格链接")

        app_token = url_info["app_token"]
        table_id = url_info.get("table_id", "")
        if not table_id:
            raise ValueError("链接中未包含 table_id")

        # 获取所有记录
        resp = requests.get(
            f"https://open.feishu.cn/open-apis/bitable/v1/apps/{app_token}/tables/{table_id}/records?page_size=500",
            headers=self._headers())
        data = resp.json()
        if data.get("code") != 0:
            raise RuntimeError(f"获取多维表格记录失败: {data}")

        results = []
        for i, record in enumerate(data["data"].get("items", [])):
            fields = record.get("fields", {})
            content = ""
            if column_name in fields:
                val = fields[column_name]
                if isinstance(val, list):
                    content = "".join([v.get("text", "") for v in val if isinstance(v, dict)])
                elif isinstance(val, str):
                    content = val
            results.append({
                "row": i,
                "record_id": record["record_id"],
                "content": content
            })
        return results

    def write_bitable_column(self, url: str, column_name: str, data: List[Dict],
                             batch_delay: float = 0.3) -> Dict:
        """向多维表格指定列写入数据"""
        url_info = self.parse_url(url)
        if url_info["type"] != "bitable":
            raise ValueError("仅支持多维表格链接")

        app_token = url_info["app_token"]
        table_id = url_info.get("table_id", "")

        stats = {"success": 0, "failed": 0}
        for item in data:
            record_id = item.get("record_id")
            text = item["content"]
            if not record_id:
                stats["failed"] += 1
                continue

            payload = {"fields": {column_name: text}}
            resp = requests.put(
                f"https://open.feishu.cn/open-apis/bitable/v1/apps/{app_token}/tables/{table_id}/records/{record_id}",
                headers=self._headers(), json=payload)
            if resp.json().get("code") == 0:
                stats["success"] += 1
            else:
                stats["failed"] += 1
            time.sleep(batch_delay)

        return stats


# ============================================================
# CLI 入口（独立使用时）
# ============================================================
def main():
    import argparse
    parser = argparse.ArgumentParser(description="飞书文档通用读写工具")
    sub = parser.add_subparsers(dest="command")

    # read 子命令
    read_p = sub.add_parser("read", help="读取飞书表格列")
    read_p.add_argument("--url", required=True, help="飞书文档链接")
    read_p.add_argument("--column", required=True, help="目标列名称")
    read_p.add_argument("--table-index", type=int, default=0)
    read_p.add_argument("--start-row", type=int, default=1)
    read_p.add_argument("--end-row", type=int, default=None)
    read_p.add_argument("--output", help="输出为JSON文件")

    # write 子命令
    write_p = sub.add_parser("write", help="写入飞书表格列")
    write_p.add_argument("--url", required=True, help="飞书文档链接")
    write_p.add_argument("--column", required=True, help="目标列名称")
    write_p.add_argument("--input", required=True, help="输入JSON文件")
    write_p.add_argument("--table-index", type=int, default=0)
    write_p.add_argument("--skip-existing", action="store_true", default=True)

    # info 子命令
    info_p = sub.add_parser("info", help="查看飞书文档表格结构")
    info_p.add_argument("--url", required=True, help="飞书文档链接")

    args = parser.parse_args()
    io = FeishuDocIO()

    if args.command == "read":
        url_info = io.parse_url(args.url)
        if url_info["type"] == "bitable":
            rows = io.read_bitable_column(args.url, args.column)
        else:
            rows = io.read_table_column(args.url, args.column,
                                        args.table_index, args.start_row, args.end_row)
        print(f"读取到 {len(rows)} 行数据")
        for r in rows:
            print(f"  行{r['row']}: {r['content'][:60]}{'...' if len(r['content'])>60 else ''}")
        if args.output:
            with open(args.output, "w", encoding="utf-8") as f:
                json.dump(rows, f, ensure_ascii=False, indent=2)
            print(f"已保存到 {args.output}")

    elif args.command == "write":
        with open(args.input, "r", encoding="utf-8") as f:
            data = json.load(f)
        url_info = io.parse_url(args.url)
        if url_info["type"] == "bitable":
            stats = io.write_bitable_column(args.url, args.column, data)
        else:
            stats = io.write_table_column(args.url, args.column, data,
                                          args.table_index, args.skip_existing)
        print(f"\n写入完成: 成功{stats['success']} 跳过{stats.get('skipped',0)} 失败{stats['failed']}")

    elif args.command == "info":
        url_info = io.parse_url(args.url)
        print(f"链接类型: {url_info['type']}")
        if url_info["type"] in ["docx", "wiki"]:
            if url_info["type"] == "wiki":
                doc_id = io._resolve_wiki_to_doc_id(url_info["wiki_token"])
                print(f"Wiki token: {url_info['wiki_token']} → doc_id: {doc_id}")
            else:
                doc_id = url_info["doc_id"]
            tables = io._find_tables(doc_id)
            print(f"找到 {len(tables)} 个表格:")
            for i, t in enumerate(tables):
                print(f"  表格{i}: {t['rows']}行 x {t['cols']}列 (block_id={t['block_id']})")
    else:
        parser.print_help()


if __name__ == "__main__":
    main()
