#!/usr/bin/env python3
"""
测试飞书云文档图片上传
"""

import requests
import os
from datetime import datetime, timedelta
from PIL import Image, ImageDraw

# ========== 配置 ==========
ZHIYUN_APP_ID = "cli_a9f29dca82b9dbef"
ZHIYUN_APP_SECRET = "sDfhjG7QT1S4gfHiMVYSygmPQPN1R2Ho"
BASE_URL = "https://open.feishu.cn/open-apis"


class FeishuDocx:
    """飞书云文档操作工具类"""

    def __init__(self):
        self._token = None
        self._token_expires = None

    @property
    def token(self) -> str:
        if self._token and self._token_expires and datetime.now() < self._token_expires:
            return self._token

        url = f"{BASE_URL}/auth/v3/tenant_access_token/internal"
        response = requests.post(url, json={
            "app_id": ZHIYUN_APP_ID,
            "app_secret": ZHIYUN_APP_SECRET
        })
        data = response.json()

        if data.get("code") != 0:
            raise Exception(f"获取 token 失败: {data}")

        self._token = data["tenant_access_token"]
        self._token_expires = datetime.now() + timedelta(seconds=data.get("expire", 7200) - 60)
        print(f"[OK] Token 获取成功")
        return self._token

    @property
    def headers(self):
        return {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json"
        }

    def set_document_permission(self, document_id: str, editable: bool = True) -> bool:
        """设置文档权限为组织内可编辑"""
        url = f"{BASE_URL}/drive/v1/permissions/{document_id}/public"

        payload = {
            "external_access_entity": "open",
            "security_entity": "anyone_can_view",
            "comment_entity": "anyone_can_view",
            "share_entity": "anyone",
            "link_share_entity": "tenant_editable" if editable else "tenant_readable",
            "invite_external": False
        }

        response = requests.patch(url, headers=self.headers, params={"type": "docx"}, json=payload)
        result = response.json()
        if result.get("code") == 0:
            print(f"[OK] 权限设置成功: 组织内可编辑")
            return True
        else:
            print(f"[WARN] 权限设置失败: {result.get('msg')}")
            return False

    def create_document(self, title: str, folder_token: str = None, editable: bool = True) -> dict:
        """创建云文档（自动设置为组织内可编辑）"""
        url = f"{BASE_URL}/docx/v1/documents"
        payload = {"title": title}
        if folder_token:
            payload["folder_token"] = folder_token

        response = requests.post(url, headers=self.headers, json=payload)
        data = response.json()

        if data.get("code") != 0:
            raise Exception(f"创建文档失败: {data}")

        doc = data["data"]["document"]
        print(f"[OK] 文档创建成功: {doc['document_id']}")

        # 自动设置权限
        if editable:
            self.set_document_permission(doc['document_id'], editable=True)

        return doc

    def create_empty_image_block(self, document_id: str, index: int = -1) -> str:
        """
        创建空的图片块，返回 block_id

        正确流程第一步：先创建空图片块
        """
        url = f"{BASE_URL}/docx/v1/documents/{document_id}/blocks/{document_id}/children"

        # 创建空的图片块
        image_block = {
            "block_type": 27,
            "image": {}  # 空的图片块
        }

        payload = {"children": [image_block]}
        if index >= 0:
            payload["index"] = index

        response = requests.post(url, headers=self.headers, json=payload)
        data = response.json()

        if data.get("code") != 0:
            raise Exception(f"创建图片块失败: {data}")

        # 获取创建的图片块 ID
        children = data["data"].get("children", [])
        if not children:
            raise Exception("创建图片块失败: 没有返回 children")

        block_id = children[0].get("block_id")
        print(f"[OK] 空图片块创建成功, block_id: {block_id}")
        return block_id

    def upload_image_to_block(self, file_path: str, block_id: str) -> str:
        """
        上传图片到指定的图片块

        正确流程第二步：将图片上传并绑定到已创建的图片块
        """
        url = f"{BASE_URL}/drive/v1/medias/upload_all"

        file_name = os.path.basename(file_path)
        file_size = os.path.getsize(file_path)

        headers = {"Authorization": f"Bearer {self.token}"}

        with open(file_path, 'rb') as f:
            files = {'file': (file_name, f, 'image/png')}
            data = {
                'file_name': file_name,
                'parent_type': 'docx_image',  # 云文档图片
                'parent_node': block_id,       # 关键: 使用图片块的 block_id
                'size': str(file_size)
            }

            print(f"[INFO] 上传图片到 block_id={block_id}")
            response = requests.post(url, headers=headers, files=files, data=data)

        result = response.json()
        print(f"[DEBUG] 上传响应: code={result.get('code')}, msg={result.get('msg')}")

        if result.get("code") != 0:
            raise Exception(f"上传图片失败: {result}")

        file_token = result["data"]["file_token"]
        print(f"[OK] 图片上传成功, file_token: {file_token}")
        return file_token

    def _upload_media(self, file_path: str, parent_type: str, parent_node: str = '') -> str:
        """使用 media API 上传"""
        url = f"{BASE_URL}/drive/v1/medias/upload_all"

        file_name = os.path.basename(file_path)
        file_size = os.path.getsize(file_path)
        headers = {"Authorization": f"Bearer {self.token}"}

        with open(file_path, 'rb') as f:
            files = {'file': (file_name, f, 'image/png')}
            data = {
                'file_name': file_name,
                'parent_type': parent_type,
                'parent_node': parent_node,
                'size': str(file_size)
            }

            response = requests.post(url, headers=headers, files=files, data=data)

        result = response.json()
        print(f"[DEBUG] media 响应 ({parent_type}): code={result.get('code')}, msg={result.get('msg')}")

        if result.get("code") == 0:
            return result["data"]["file_token"]
        return None

    def _upload_to_drive(self, file_path: str) -> str:
        """上传到云空间根目录，然后获取 file_token"""
        # 先获取根文件夹 token
        url = f"{BASE_URL}/drive/explorer/v2/root_folder/meta"
        headers = {"Authorization": f"Bearer {self.token}"}

        response = requests.get(url, headers=headers)
        result = response.json()
        print(f"[DEBUG] 根目录响应: {result}")

        if result.get("code") != 0:
            return None

        root_token = result["data"]["token"]

        # 上传文件
        file_name = os.path.basename(file_path)
        file_size = os.path.getsize(file_path)

        upload_url = f"{BASE_URL}/drive/v1/medias/upload_all"
        with open(file_path, 'rb') as f:
            files = {'file': (file_name, f, 'image/png')}
            data = {
                'file_name': file_name,
                'parent_type': 'explorer',
                'parent_node': root_token,
                'size': str(file_size)
            }

            response = requests.post(upload_url, headers=headers, files=files, data=data)

        result = response.json()
        print(f"[DEBUG] drive 上传响应: code={result.get('code')}, msg={result.get('msg')}")

        if result.get("code") == 0:
            return result["data"]["file_token"]
        return None

    def _upload_im_image(self, file_path: str) -> str:
        """使用消息图片 API 上传"""
        url = f"{BASE_URL}/im/v1/images"

        headers = {"Authorization": f"Bearer {self.token}"}

        with open(file_path, 'rb') as f:
            files = {'image': f}
            data = {'image_type': 'message'}

            response = requests.post(url, headers=headers, files=files, data=data)

        result = response.json()
        print(f"[DEBUG] im 图片响应: code={result.get('code')}, msg={result.get('msg')}")

        if result.get("code") == 0:
            return result["data"]["image_key"]
        return None

    def _try_upload_with_type(self, file_path: str, parent_type: str) -> str:
        """尝试不同的 parent_type 上传"""
        url = f"{BASE_URL}/drive/v1/medias/upload_all"

        file_name = os.path.basename(file_path)
        file_size = os.path.getsize(file_path)

        headers = {"Authorization": f"Bearer {self.token}"}

        with open(file_path, 'rb') as f:
            files = {'file': (file_name, f, 'image/png')}
            data = {
                'file_name': file_name,
                'parent_type': parent_type,
                'parent_node': '',
                'size': str(file_size)
            }

            print(f"[INFO] 尝试 parent_type={parent_type}")
            response = requests.post(url, headers=headers, files=files, data=data)

        result = response.json()
        print(f"[DEBUG] 响应: {result}")

        if result.get("code") != 0:
            raise Exception(f"上传失败 ({parent_type}): {result}")

        return result["data"]["file_token"]

    def bind_image_to_block(self, document_id: str, block_id: str, file_token: str) -> dict:
        """
        绑定图片到图片块 (关键的第三步!)

        使用 PATCH 请求和 replace_image 字段将图片绑定到已创建的图片块
        """
        url = f"{BASE_URL}/docx/v1/documents/{document_id}/blocks/{block_id}"

        params = {"document_revision_id": -1}

        payload = {
            "replace_image": {
                "token": file_token
            }
        }

        print(f"[INFO] 绑定图片到块: block_id={block_id}")
        response = requests.patch(url, headers=self.headers, params=params, json=payload)
        data = response.json()

        if data.get("code") != 0:
            raise Exception(f"绑定图片失败: {data}")

        print(f"[OK] 图片绑定成功")
        return data["data"]


def generate_test_image(output_path: str):
    """生成测试图片"""
    img = Image.new('RGB', (400, 300), color='#4a90d9')
    draw = ImageDraw.Draw(img)
    draw.rectangle([20, 20, 380, 280], fill='#f0f4f8', outline='#2563eb', width=2)
    draw.text((200, 150), "Test Image", fill='#1e3a5f', anchor='mm')
    draw.text((200, 200), datetime.now().strftime("%Y-%m-%d %H:%M:%S"), fill='#6b7280', anchor='mm')
    img.save(output_path)
    print(f"[OK] 测试图片生成: {output_path}")
    return output_path


def main():
    print("\n" + "#" * 60)
    print("#  飞书云文档图片上传测试")
    print("#" * 60)

    docx = FeishuDocx()

    # Step 1: 生成测试图片
    print("\n" + "=" * 50)
    print("Step 1: 生成测试图片")
    print("=" * 50)

    test_image = "/tmp/feishu_docx_test_image.png"
    generate_test_image(test_image)

    # Step 2: 创建测试文档
    print("\n" + "=" * 50)
    print("Step 2: 创建测试文档")
    print("=" * 50)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    doc = docx.create_document(f"图片上传测试 - {timestamp}")
    document_id = doc["document_id"]

    # Step 3: 创建空的图片块
    print("\n" + "=" * 50)
    print("Step 3: 创建空的图片块")
    print("=" * 50)

    try:
        block_id = docx.create_empty_image_block(document_id)
    except Exception as e:
        print(f"[ERROR] 创建图片块失败: {e}")
        return

    # Step 4: 上传图片
    print("\n" + "=" * 50)
    print("Step 4: 上传图片")
    print("=" * 50)

    try:
        file_token = docx.upload_image_to_block(test_image, block_id)
    except Exception as e:
        print(f"[ERROR] 上传图片失败: {e}")
        return

    # Step 5: 绑定图片到图片块 (关键步骤!)
    print("\n" + "=" * 50)
    print("Step 5: 绑定图片到图片块")
    print("=" * 50)

    try:
        docx.bind_image_to_block(document_id, block_id, file_token)
    except Exception as e:
        print(f"[ERROR] 绑定图片失败: {e}")
        return

    # 完成
    print("\n" + "=" * 60)
    print("测试完成!")
    print("=" * 60)
    print(f"\n文档地址:")
    print(f"  https://feishu.cn/docx/{document_id}")
    print()


def check_permissions(docx):
    """检查应用权限"""
    # 获取应用信息
    url = f"{BASE_URL}/application/v6/applications/underauditlist"
    headers = {"Authorization": f"Bearer {docx.token}"}

    response = requests.get(url, headers=headers)
    print(f"[DEBUG] 权限检查响应: {response.json()}")


if __name__ == "__main__":
    main()