#!/usr/bin/env python3
"""
上传用户指定的图片到飞书云文档
"""

import requests
import os
from datetime import datetime, timedelta

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

# 用户指定的图片
IMAGE_PATH = "/Users/donglinlai/Downloads/u274.png"

_token = None
_token_expires = None


def get_token():
    global _token, _token_expires
    if _token and _token_expires and datetime.now() < _token_expires:
        return _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}")

    _token = data["tenant_access_token"]
    _token_expires = datetime.now() + timedelta(seconds=data.get("expire", 7200) - 60)
    return _token


def headers():
    return {
        "Authorization": f"Bearer {get_token()}",
        "Content-Type": "application/json"
    }


def set_document_permission(document_id: str, editable: bool = True):
    """
    设置文档权限

    Args:
        document_id: 文档ID
        editable: True=组织内可编辑, False=组织内只读
    """
    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=headers(), params={"type": "docx"}, json=payload)
    data = response.json()

    if data.get("code") == 0:
        print(f"      权限设置成功: {'组织内可编辑' if editable else '组织内只读'}")
        return True
    else:
        print(f"      [WARN] 权限设置失败: {data.get('msg')}")
        return False


def create_document(title: str, editable: bool = True):
    """创建文档并设置权限"""
    url = f"{BASE_URL}/docx/v1/documents"
    response = requests.post(url, headers=headers(), json={"title": title})
    data = response.json()
    if data.get("code") != 0:
        raise Exception(f"创建文档失败: {data}")

    doc = data["data"]["document"]
    document_id = doc["document_id"]

    # 自动设置权限
    if editable:
        set_document_permission(document_id, editable=True)

    return document_id


def create_image_block(document_id: str):
    """创建空图片块"""
    url = f"{BASE_URL}/docx/v1/documents/{document_id}/blocks/{document_id}/children"
    payload = {
        "children": [{
            "block_type": 27,
            "image": {}
        }]
    }
    response = requests.post(url, headers=headers(), json=payload)
    data = response.json()
    if data.get("code") != 0:
        raise Exception(f"创建图片块失败: {data}")
    return data["data"]["children"][0]["block_id"]


def upload_image(file_path: str, block_id: str):
    """上传图片"""
    url = f"{BASE_URL}/drive/v1/medias/upload_all"

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

    # 根据扩展名设置 MIME 类型
    ext = file_name.lower().split('.')[-1]
    mime_types = {
        'png': 'image/png',
        'jpg': 'image/jpeg',
        'jpeg': 'image/jpeg',
        'gif': 'image/gif',
        'webp': 'image/webp'
    }
    mime_type = mime_types.get(ext, 'image/png')

    with open(file_path, 'rb') as f:
        files = {'file': (file_name, f, mime_type)}
        data = {
            'file_name': file_name,
            'parent_type': 'docx_image',
            'parent_node': block_id,
            'size': str(file_size)
        }

        response = requests.post(
            url,
            headers={"Authorization": f"Bearer {get_token()}"},
            files=files,
            data=data
        )

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

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


def bind_image(document_id: str, block_id: str, file_token: str):
    """绑定图片到图片块"""
    url = f"{BASE_URL}/docx/v1/documents/{document_id}/blocks/{block_id}"
    payload = {
        "replace_image": {
            "token": file_token
        }
    }
    response = requests.patch(
        url,
        headers=headers(),
        params={"document_revision_id": -1},
        json=payload
    )
    data = response.json()
    if data.get("code") != 0:
        raise Exception(f"绑定失败: {data}")
    return True


def main():
    print(f"\n上传图片: {IMAGE_PATH}")
    print("=" * 60)

    # 检查文件
    if not os.path.exists(IMAGE_PATH):
        print(f"[ERROR] 文件不存在: {IMAGE_PATH}")
        return

    file_size = os.path.getsize(IMAGE_PATH)
    print(f"文件大小: {file_size / 1024:.1f} KB")

    # Step 1: 创建文档
    print("\n[1/4] 创建飞书文档...")
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M")
    doc_id = create_document(f"火把图片 - {timestamp}")
    print(f"      文档ID: {doc_id}")

    # Step 2: 创建图片块
    print("[2/4] 创建图片块...")
    block_id = create_image_block(doc_id)
    print(f"      块ID: {block_id}")

    # Step 3: 上传图片
    print("[3/4] 上传图片...")
    file_token = upload_image(IMAGE_PATH, block_id)
    print(f"      file_token: {file_token}")

    # Step 4: 绑定图片
    print("[4/4] 绑定图片...")
    bind_image(doc_id, block_id, file_token)
    print("      绑定成功!")

    print("\n" + "=" * 60)
    print("上传完成!")
    print(f"\n文档地址: https://feishu.cn/docx/{doc_id}")
    print("=" * 60)


if __name__ == "__main__":
    main()